from enum import Enum
from typing import Any, Literal, Optional
from pydantic import BaseModel, Field

from app.configs.settings import settings
from app.services.pinecone import retrieve
from langchain_openai.chat_models.base import ChatOpenAI
from langchain_core.output_parsers.json import parse_partial_json
from langchain_core.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate, AIMessagePromptTemplate
from app.utils.constants import BASE_SYSTEM_TEMPLATE, get_hotel_details, SYSTEM_TEMPLATE_FOR_CONTEXT_DATA

class ToolEnum(Enum):
    modifyReservation = 'modify_reservation_tool'
    search = "search_room_tool"
    booking = "add_to_cart"
    billing = "add_billing_details"
    send_email = "send_email_tool"
    clear_cart = "clear_cart_tool"
    special_request = "note_special_request"
    get_reservation_details = "get_reservation_details"
    default = "no_tool"

class AgentParams(BaseModel):
    check_in: str = Field(
        description=(
            "Check-in date in MM-DD-YYYY format. "
            "Always populate this field. Use the user's check-in date if provided; otherwise use the default CHECK-IN value from the prompt."
        )
    )
    check_out: str = Field(
        description=(
            "Check-out date in MM-DD-YYYY format. "
            "Always populate this field. Use the user's check-out date if provided; otherwise use the default CHECK-OUT value from the prompt."
        )
    )
    adultCount: int = Field(
        default=2,
        ge=1,
        description=(
            "Number of adults for the room search or booking. "
            "Always populate this field. Use the user's value if provided; otherwise default to 2."
        )
    )
    childCount: int = Field(
        default=0,
        ge=0,
        description=(
            "Number of children for the room search or booking. "
            "Always populate this field. Use the user's value if provided; otherwise default to 0."
        )
    )
    room_name: list[str] = Field(
        default=[],
        description="Requested room names if the user mentions one explicitly."
    )
    room_type_id: list[str] = Field(
        default=[],
        description=(
            "List of exact room_type_ids for booking or add-to-cart requests. "
            "Use only when the requested rooms is known from context or prior results. "
            "Never guess or fabricate this value."
        )
    )
    confirmation_number: str = Field(
        default=None,
        description=(
            "Reservation confirmation number for cancellation-email requests. Itenrary is also accepted here. "
            "Reservation confirmation number for special request notes if reservation is already confirmed. Itenrary is also accepted here."
            "Reservation confirmation number for getting reservation details when calling get_reservation_details tool. Itenrary is also accepted here."
            "Use a value from the user query or conversation history when available."
        )
    )
    email: str = Field(
        default=None,
        description=(
            "Email address of the user for special request notes. "
            "Email address of the user for getting reservation details when calling get_reservation_details tool. "
            "Use a value from the user query or conversation history when available."
        )
    )
    special_request: str = Field(
        default=None,
        description=(
            "Special request notes from the user for the reservation. "
            "Don't copy exact text from user query, Explain the request in your own words and summarize it in your own words."
        )
    )

class SortDirection(str, Enum):
    ascending = "ascending"
    descending = "descending"

class SortField(str, Enum):
    rate = "rate"
    occupancy = "occupancy"

class FilterField(str, Enum):
    rate = "rate"
    occupancy = "occupancy"

class FilterOperator(str, Enum):
    eq = "=="       # equal
    ne = "!="       # not equal
    gt = ">"        # greater than
    lt = "<"        # less than
    gte = ">="      # greater than equal
    lte = "<="      # less than equal

class FilterCondition(BaseModel):
    field: FilterField = Field(description="Field name to filter on (e.g., rate, occupancy)")
    operator: FilterOperator = Field(description="Comparison operator")
    value: Any = Field(description="Value to compare against")


class SortConfig(BaseModel):
    field: SortField = Field(description="Field to sort by")
    direction: SortDirection = Field(description="Sort direction")

class Filters(BaseModel):
    sort: SortConfig | None = Field(
        default=None,
        description="Sorting configuration"
    )
    conditions: list[FilterCondition] = Field(
        default_factory=list,
        description="List of filter conditions"
    )

class BillingFormStep(str, Enum):
    addons = "addons"
    insurance = "insurance"
    customer_details_form = "customer_details_form"
    promo_code_form = "promo_code_form"
    payment_details_form = "payment_details_form"

class AgentSchema(BaseModel):
    tool: ToolEnum = Field(
        description=(
            "Tool to execute for the user's request. "
            "Use add_to_cart for booking or add-to-cart actions, "
            "search_room_tool for room search or reopening the room slider, "
            "send_email_tool for cancellation email requests, "
            "note_special_request for special request notes, "
            "and no_tool when only a conversational answer is needed. "
            "The no_tool path may still include room_type_ids so the frontend can open the slider from general room matches."
        )
    )
    billing_form_step: Optional[BillingFormStep] = Field(
        default=None,
        description="if calling the billing tool, the step of the billing form to show to the user"
    )
    answer: str = Field(
        description=(
            "Friendly user-facing response. Keep it positive and concise. "
            "If a tool will be used, start naturally with 'Let me...'."
        )
    )
    params: AgentParams = Field(
        description="Structured parameters extracted from the user's request."
    )
    filters: Filters = Field(
        default=None,
        description="Filters to apply to the room search"
    )
    room_count: int = Field(
        default=-1,
        description="Number of rooms to return in the frontend slider. Use -1 to return all matching rooms."
    )
    hotel_ids: list[Literal["123", "124", "125"]] = Field(description=(
        "Ordered list of hotel IDs whose rooms should be returned.\n\n"
        "Rules:\n"
        "1. Default to ['124'] and show Bayview Plaza rooms first.\n"
        "2. If the user asks for more options or other hotels, use ['123', '124', '125'].\n"
        "3. If the user explicitly requests a specific hotel, include only the requested hotel IDs.\n"
        "4. Do not change hotels unless the user asks or signals dissatisfaction.\n"
        "5. room_type_ids must stay scoped to these hotel_ids."
    ), default_factory=lambda: ["124"])
    room_type_ids: list[str] = Field(
        default_factory=list,
        description=(
            "Specific room_type_id values to render on the frontend when search_room_tool is not being used. "
            "If the answer lists or recommends room types from GENERAL-ROOM-DETAILS, always include the matching room_type_ids here. "
            "This is required for frontend slider rendering. "
            "Only include room_type_ids from the currently selected hotel_ids. "
            "By default that means only hotel 124 unless the user explicitly asks for another hotel or broader hotel options. "
            "Leave empty only when no specific room types are being referenced."
        )
    )
    total_pets: int = Field(
        default=0,
        description="Total number of allowed pets (dogs and cats) in the room. Use 0 for no pets. Always include this field if the user explicitly asks for pets or mentions they are bringing pets. Frontend will use this field to automatically add pets add-on to the cart if user add a room to cart in the same chat. Wait for clarification from the user if you are not sure about the number of pets. If user mention any pet so in every conversation you have keep track of the number of pets and update this field accordingly becasue this is the only way to know the number of pets in the room until chat is over. Make it back to 0 if user mention no pets or no more pets or switch to hotel which doesn't allow pets."
    )

class ContextData(BaseModel):
    query: str = Field(description="The standalone query to retrieve the most relevant context for answering the user's latest question.", default="")
    skip: bool = Field(description="Whether to skip vector database retrieval for the user's latest question.", default=False)

def get_chat_prompt(history:list = [], tool_response: str | None = None, tool_call: str | None = None) -> ChatPromptTemplate:

    if not BASE_SYSTEM_TEMPLATE:
        raise ValueError("Invalid BASE_SYSTEM_TEMPLATE")

    messages = [
        SystemMessagePromptTemplate.from_template(BASE_SYSTEM_TEMPLATE),
        *history,
        HumanMessagePromptTemplate.from_template("{question}")
    ]
    if tool_response and tool_call:
        messages.append(AIMessagePromptTemplate.from_template("{tool_call}"))
        messages.append(AIMessagePromptTemplate.from_template("Tool response: {tool_response}"))
    return ChatPromptTemplate.from_messages(messages)

async def get_context_data(question: str, sub_hotel_id: str, history: list = [], callbacks: Any = None) -> dict:

    llm = ChatOpenAI(api_key=settings.OPENAI_API_KEY, model="gpt-5.4-mini")

    messages = [
        SystemMessagePromptTemplate.from_template(SYSTEM_TEMPLATE_FOR_CONTEXT_DATA),
    ]
    chain = ChatPromptTemplate.from_messages(messages) |  llm.with_structured_output(ContextData)
    result: ContextData = await chain.ainvoke({"history": history, "question": question}, config={"run_name": "Context Data", "callbacks": callbacks})

    if result.skip:
        return ""

    return await retrieve(result.query, sub_hotel_id)

class HotelAgent:

    def __init__(self, sub_hotel_id: str, model: str = "gpt-5.1", api_key: str = settings.OPENAI_API_KEY, conversation_id: str | None = None):
        
        self.sub_hotel_id = sub_hotel_id
        self.model = model
        self.api_key = api_key
        self.conversation_id: str | None = conversation_id
        self.result: AgentSchema | None = None

    async def stream(self, input: dict = {}, history: list = [], tool_response: str | None = None, tool_call: str | None = None, callbacks: Any = None):
        llm = ChatOpenAI(api_key=self.api_key, model=self.model, streaming=True)

        schema = AgentSchema.model_json_schema()
        schema["properties"]["hotel_ids"]["description"] = schema["properties"]["hotel_ids"]["description"].replace("['124']", f"['{self.sub_hotel_id}']")
        schema["properties"]["room_type_ids"]["description"] = schema["properties"]["room_type_ids"]["description"].replace("124", self.sub_hotel_id)

        llm_with_schema = llm.with_structured_output(schema)
        context = await get_context_data(input.get('question'), self.sub_hotel_id, history, callbacks)
        input['hotel_context'] = context
        chain = get_chat_prompt(history=history, tool_response=tool_response, tool_call=tool_call) | llm_with_schema
        
        if tool_response and tool_call:
            input['tool_response'] = tool_response
            input['tool_call'] = tool_call

        stream_result = chain.astream_events(
            input=input,
            config={
                "run_name": "Main Agent", 
                "tags": [self.conversation_id if self.conversation_id else 'unknown-conversation', self.sub_hotel_id if self.sub_hotel_id else 'unknown-hotel'],
                "callbacks": callbacks,
            }
        )

        final_answer = ""
        answer = ""

        async for event in stream_result:
            if event and event.get("event") == "on_chat_model_stream":
                if event.get("data") and event.get("data").get("chunk"):
                    if hasattr(event.get("data").get("chunk"), "content"):
                        token = event.get("data").get("chunk").content
                        if not token:
                            continue
                        final_answer += token
                        escaped_text = parse_partial_json(final_answer)
                        
                        if escaped_text is not None and 'answer' in escaped_text:
                            new_token = escaped_text['answer'][len(answer):]
                            answer += new_token
                            if new_token == "": 
                                continue
                            yield new_token
                        answer = answer
            
            if event and event.get("event") == "on_chain_end":
                output = event.get("data", {}).get("output")
                if output is not None:
                    self.result = AgentSchema.model_validate(output)