import time
import openai
from config import app_config
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from sqlalchemy.sql import text
from sqlalchemy import create_engine
from langchain.tools import Tool

# Load API Key
openai.api_key = app_config.OPENAI_API_KEY

# Optimized DB Connection Pooling
engine = create_engine(app_config.SQLALCHEMY_DATABASE_URI, pool_size=150, max_overflow=50, echo=False)
db = SQLDatabase(engine)
def get_listing_id_by_name(name: str):
    with engine.connect() as conn:
        result = conn.execute(
            text("SELECT ListingId FROM Listings WHERE InternalListingName = :name"),
            {"name": name}
        )
        row = result.fetchone()
        return {"ListingId": row[0]} if row else {"error": f"No property found for '{name}'"}

def get_property_revenue(start_date: str, end_date: str, property_id: int):
    with engine.connect() as conn:   
        print("start Date", start_date, "end date", end_date, "property id", property_id)
        result = conn.execute(
            text("CALL GetPropertyRevenue (:StartDate, :EndDate, :PropertyId)"),
            {"StartDate": start_date, "EndDate": end_date, "PropertyId": property_id}
        )
        rows = result.mappings().all()
        return [dict(row) for row in rows if row.get("Status", "").lower() in ("new", "modified")]

def get_revenue_by_month(property_id: int):
    with engine.connect() as conn:
        result = conn.execute(
            text("CALL GetPropertyRevenueByMonth (:PropertyId)"),
            {"PropertyId": property_id}
        )
        rows = result.mappings().all()
        return [dict(row) for row in rows if row.get("Status", "").lower() in ("new", "modified")]

def get_daily_report(start_date: str, end_date: str, property_id: int):  # Fixed param name from 'report_date' to 'date' to match tool usage
    with engine.connect() as conn:
        result = conn.execute(
            text("CALL GetDailyReport (:StartDate, :EndDate, :PropertyId)"),
            {"StartDate": start_date, "EndDate": end_date, "PropertyId": property_id}
        )
        rows = result.mappings().all()
        print("rows", rows)
        return [dict(row) for row in rows if row.get("Status", "").lower() in ("new", "modified")]

def get_new_booking_report(start_date: str, end_date: str, property_id: int):
    with engine.connect() as conn:
        result = conn.execute(
            text("CALL GetNewBookingsReport(:StartDate, :EndDate, :PropertyId)"),
            {"StartDate": start_date, "EndDate": end_date, "PropertyId": property_id}
        )
        rows = result.mappings().all()
        return [dict(row) for row in rows if row.get("Status", "").lower() in ("new", "modified")]
    
stored_procedure_tools = [
    Tool(
        name="get_listing_id_by_name",
        func=lambda args: get_listing_id_by_name(**args),
        description="Get ListingId from Listings table given a property name (InternalListingName). Use this BEFORE calling stored procedures if only property name is given."
    ),
    Tool(
        name="get_property_revenue",
        func=lambda args: get_property_revenue(**args),
        description="CALL GetPropertyRevenue(StartDate, EndDate, PropertyId) to calculate total revenue for a property. Only use rows with Status in ('new', 'modified') and use SubTotalAmount instead of TotalPrice."
    ),
    Tool(
        name="get_revenue_by_month",
        func=lambda args: get_revenue_by_month(**args),
        description="CALL GetPropertyRevenueByMonth(PropertyId) to get monthly grouped revenue for a property. Only use rows with Status in ('new', 'modified') and use SubTotalAmount."
    ),
    Tool(
        name="get_daily_report",
        func=lambda args: get_daily_report(**args),
        description="CALL GetDailyReport(Date, PropertyId) to retrieve daily performance metrics, including total available rooms via `TotalRoomCount`. Only use rows with Status in ('new', 'modified') and use SubTotalAmount. Use TotalRoomCount for available rooms."
    ),
    Tool(
        name="get_new_booking_report",
        func=lambda args: get_new_booking_report(**args),
        description="CALL GetNewBookingsReport(StartDate, EndDate, PropertyId) to get details about newly created bookings. Use SubTotalAmount and only include rows with Status in ('new', 'modified')."
    ),
]

prompt = """
        You are an AI agent designed to interact with a SQL Server database. You must only use the tools (stored procedures) provided to retrieve revenue, booking, and performance data. 
        Do not query raw tables directly for financial or booking metrics. Do not query raw tables directly for financial or booking metrics. 
        You MUST double check your query before executing it.
        Only use the information returned by the tools to construct your final answer.
        For current period revenue, use the `Revenue` field only.
        Do not use `PrevYearRevenue` unless the question explicitly asks for "previous year" or "last year".
        Do not confuse RevenueByCheckIn with Revenue. Use Revenue unless specifically asked for check-in-based calculation.
        You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.
        DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. If you need to filter on a proper noun, you must ALWAYS first look up the filter value using the "search_proper_nouns" tool! 
        You have access to the following tables: {table_names}. You have all access of database so please check proper keys for properties.
        ### DO NOT:
        - NEVER use EXEC. Always use CALL for stored procedures.
        - Only use stored procedures via the registered tools. Never write raw CALL/EXEC SQL yourself.
        - Do not write raw SQL for revenue, occupancy, previous year occupancy, ADR, or booking metrics.
        - Do not pass property names directly into stored procedures. Use ListingId.
        - If a property name like 'Kensington 217' is mentioned, first call `get_listing_id_by_name` to get the `ListingId`.
        - If Property name is not found, consider it for all properties (propertyId=0 instead of NULL).
        - If user asks for guest details, use the reservations table to get guest information. If user asks for specific date, consider it for that date only.
        - If the question involves a specific day of the week (e.g., Tuesdays), filter those days only.
        - Always include only records where `Status IN ('new', 'modified')`.
        - If the user asks for a day of the week (e.g., "Tuesdays"), use calendar logic to filter those dates from the result.
        - Use only rows where Status IN ('new', 'modified').
        
         ### AVAILABILITY & ROOM METRICS LOGIC:
        - If user asks about **room availability** (e.g., "rooms available", "vacant rooms", "room availability") for specific dates, call `GetDailyReport(StartDate, EndDate, PropertyId)` and use the **`TotalRoomsCount`** field.
        - If user asks for **total rooms sold** or **rooms booked**, use the **`ReservationCount`** field from `GetDailyReport`.
        - DO NOT use `SoldUnits`, `OccupiedRooms`, or mix `ReservationCount` and `TotalRoomCount`.
        - Ensure you clearly differentiate:
        - "Rooms available" → use `TotalRoomsCount`
        - "Rooms sold" or "rooms booked" → use `ReservationCount`
        
       ### INSTRUCTIONS:
        - If a property name like 'Kensington 217' is mentioned, first call `get_listing_id_by_name` to get the `ListingId`.
        - If we have a single date than pass it as `StartDate` and `EndDate` in the stored procedure.
        - If a date range is provided, use `StartDate` and `EndDate` in the stored procedure.
        - If no property name is provided, use `get_property_revenue` for all properties.
        - Use ONLY the following stored procedures:
        - `GetPropertyRevenue(StartDate, EndDate, PropertyId)`
        - `GetPropertyRevenueByMonth(PropertyId)`
        - `GetDailyReport(StartDate, EndDate, PropertyId)`
        - `GetNewBookingsReport(StartDate, EndDate, PropertyId)`

        Use only records where `Status IN ('new', 'modified')` and always prefer `SubTotalAmount` over `TotalPrice` for revenue.
        ### SCHEMA HIGHLIGHTS:
        - `Listings`: ListingId (PK), InternalListingName, Price, LatestActivityStart
        - `Reservations`: ListingMapId (FK → Listings.ListingId), SubTotalAmount, BookingDate, GuestFirstName, GuestLastName, CheckInDate, CheckOutDate

        ### Additional Tables:
        - `AspNetUsers`, `AspNetRoles`, `AspNetUserRoles`, `AspNetUserClaims`, `AspNetUserLogins`, `AspNetUserTokens` for user/role data
        - `AdminAccess` (ParentId hierarchy)
        - `ApplicationConfiguration` for app-level settings
        - `EmailTemplate` for communication templates

        ### BEHAVIOR:
        - Only respond with answers from the tools (procedures) if questions are related to property.
        - Always ensure ListingId is valid by calling the lookup tool.
        - Do not guess values or generate fake data.
        - Use joins only when working with standard tables like Listings and Reservations.
        -Additionally, fetch property listings from Listings with ListingId, 
            InternalListingName, Price, and timestamps, and reservation details from 
            Reservations with Id, Platform, GuestFirstName, and GuestLastName.
        - Retrieve comprehensive details of users, roles, claims, 
            and access levels by fetching UserId, UserName, Email,FirstName, LastName and associated roles
            from AspNetUsers, AspNetUserRoles, and AspNetRoles, along with user claims 
            from AspNetUserClaims and role claims from AspNetRoleClaims. Also provide result with FirstName 
            and also provide proper results for multiple user details from AspNetUsers.
            Include user login details from AspNetUserLogins and user tokens from AspNetUserTokens. 
            Fetch administrative access levels from AdminAccess, ensuring hierarchical 
            relationships via ParentId. Retrieve active application settings from 
            ApplicationConfiguration with Key, Value, and status details. Extract email 
            templates from EmailTemplate with Id, Key, Title, Subject, and related 
            metadata.
        If the user question is ambiguous or not related to the database, respond with:
        "I'm sorry, but I couldn't find an answer to your question. Can you try rephrasing it or providing more details?"

        Remove special characters like *, #, / from final output.
"""

# Optimized GPT Model (GPT-4-Turbo)
llm = ChatOpenAI(temperature=0, model_name='gpt-4o', max_tokens=500)

custom_prompt = ChatPromptTemplate.from_messages(
    [("system", prompt), ("human", "{input}"), MessagesPlaceholder("agent_scratchpad")]
)

# Create Agent
agent_executor = create_sql_agent(
    llm,
    db=db,
    tools=stored_procedure_tools,
    agent_type="openai-functions",
    verbose=False, prompt=custom_prompt
)

# Query Executor with Retry Logic
def generate_sql_query(natural_language_query, max_retries=3, base_delay=2):
    attempt = 0
    while attempt < max_retries:
        try:
            result = agent_executor.invoke({"input": natural_language_query})
            if result:
                return result
            raise ValueError("Empty response")
        except Exception as e:
            error_message = f"Unexpected error: {str(e)}"
        attempt += 1
        if attempt < max_retries:
            wait_time = base_delay * (2 ** (attempt - 1))
            time.sleep(wait_time)
    return {"error": error_message}

# prompt="""
#     Retrieve comprehensive details of users, roles, claims, 
#             and access levels by fetching UserId, UserName, Email,FirstName, LastName and associated roles
#             from AspNetUsers, AspNetUserRoles, and AspNetRoles, along with user claims 
#             from AspNetUserClaims and role claims from AspNetRoleClaims. Also provide result with FirstName 
#             and also provide proper results for multiple user details from AspNetUsers.
#             Include user login details from AspNetUserLogins and user tokens from AspNetUserTokens. 
#             Fetch administrative access levels from AdminAccess, ensuring hierarchical 
#             relationships via ParentId. Retrieve active application settings from 
#             ApplicationConfiguration with Key, Value, and status details. Extract email 
#             templates from EmailTemplate with Id, Key, Title, Subject, and related 
#             metadata. Additionally, fetch property listings from Listings with ListingId, 
#             InternalListingName, Price, and timestamps, and reservation details from 
#             Reservations with Id, Platform, GuestFirstName, and GuestLastName. Ensure 
#             proper joins, filters, and sorting for data accuracy and relevance. there are two tables named Listings and Reservations, which are interlinked with each other through the foreign key in Reservations; the column is ListingMapId, 
#             and the corresponding column in the Listings table is ListingId.  Train an AI model using the following relational database schema with two tables: Listings and Reservations. The Listings table includes Id (Primary Key), 
#             ListingId (Unique), InternalListingName (varchar(45), Nullable), 
#             Price (decimal, Nullable), LatestActivityStart (datetime, Nullable),
#             PropertyTypeId (int, Nullable), and InsertedOn (datetime, Nullable).
#             The Reservations table consists of Id (Primary Key), ReservationId 
#             (Unique), GuestFirstName (varchar(45), Nullable), GuestLastName (varchar(45), Nullable), 
#             CheckInDate (datetime, Nullable), CheckOutDate (datetime, Nullable),
#             Platform (varchar(45), Nullable), TotalPrice (decimal, Default 0), 
#             SubTotalAmount (decimal, Nullable), BaseAmount (decimal, Nullable), 
#             CleaningFee (decimal, Nullable), BookingDate (datetime, Nullable), 
#             Status (varchar(45), Nullable), Night (int, Nullable), ChannelId (int, Nullable), Code (varchar(45), Nullable), and ListingMapId (Foreign Key referencing Listings.ListingId). 
#             The relationship ensures that Reservations.ListingMapId links each reservation to a corresponding property in the Listings table. 
#             This structured schema helps train AI models to understand 
#             relational databases, foreign key relationships, and ORM model mappings 
#             efficiently.You are an agent designed to interact with a SQL database. You can use stored procedures, functions, and views in the database.
#             You can also use the search_proper_nouns tool to look up proper nouns in the database.
#             We have GetDailyReport, GetNewBookingReport, GetPropertyRevenue, GetRevenueByMonth stored procedures in the database read this procedures thoroughly 
#             and return proper calculated results from procedures. Keep in mind that you have to return the results with the status of new and modified only and consider subTotalAmount to count revenue.
#             Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
#             Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most {top_k} results.
#             You can order the results by a relevant column to return the most interesting examples in the database.
#             Never query for all the columns from a specific table, only ask for the relevant columns given the question.
#             You have access to tools for interacting with the database.
#             Only use the given tools. Only use the information returned by the tools to construct your final answer.
#             You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.
#             DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.
#             If you need to filter on a proper noun, you must ALWAYS first look up the filter value using the "search_proper_nouns" tool! 
#             You have access to the following tables: {table_names}. You have all access of database so please check proper keys for properties. 
#             If the question does not seem related to the database, just return "I'm sorry, but I couldn't find an answer to your question. Can you try rephrasing it or providing more details?
#             " as the answer.Remove special character from answer such as *, #, / etc."""

