"""FoodHub Chat Agent — backend logic for the Streamlit chatbot app.

This module initialises the LLM, SQL Agent, tools, and Chat Agent,
and exposes run_chat_agent_query() as the single entry point for the
Streamlit frontend.
"""

import os
from typing import Annotated

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain.chat_models import init_chat_model
from langchain.tools import tool
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import HumanMessage, SystemMessage

# Load environment variables — OPENAI_API_KEY is read from the environment.
# Locally: set in .env file. On Hugging Face Space: set as a Space secret.
load_dotenv()

# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

# GPT-5-mini with medium reasoning — same configuration as used in the notebook.
# temperature=0.2 for deterministic, accurate responses.
model_gpt5_mini = init_chat_model(
    model="gpt-5-mini",
    model_provider="openai",
    temperature=0.2,
    max_tokens=1024,
    reasoning={"effort": "medium", "summary": "auto"},
)

# ---------------------------------------------------------------------------
# Database
# ---------------------------------------------------------------------------

# Resolve the DB path relative to this file so the path is correct whether
# the app is run locally or inside Docker (where the working directory may differ).
_DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "customer_orders.db")
db = SQLDatabase.from_uri(f"sqlite:///{_DB_PATH}")

# ---------------------------------------------------------------------------
# SQL Agent
# ---------------------------------------------------------------------------

SQL_AGENT_SYSTEM_PROMPT = f"""
You are a SQL data retrieval agent for FoodHub's order management database.
Your role is to translate customer order queries into SQL, execute them against the {db.dialect} database, and return structured results.

Database schema notes (read carefully before querying):
- All time columns (order_time, preparing_eta, prepared_time, delivery_eta, delivery_time)
  are stored as TEXT in 'HH:MM' format (e.g., '12:30').
  To compute a time difference in minutes between two HH:MM columns, use:
    round((julianday('1970-01-01 ' || col_end) - julianday('1970-01-01 ' || col_start)) * 24 * 60)
- The item_in_order column is a comma-separated TEXT string (e.g., 'Burger, Fries').
  To count the number of items in an order, use:
    (length(item_in_order) - length(replace(item_in_order, ',', '')) + 1)
- NULL values in time columns mean that stage has not been reached yet
  (e.g., delivery_time is NULL if the order has not yet been delivered).

Query rules:
1. ALWAYS inspect the table schema before writing any query.
2. Query ONLY for the specific order_id provided — never return data for all orders at once.
3. ALWAYS include cust_id in the result for customer ownership verification.
4. Limit results to at most 5 rows unless explicitly instructed otherwise.
5. Double-check your SQL before executing. If an error occurs, rewrite and retry once.
6. DO NOT execute any DML statements (INSERT, UPDATE, DELETE, DROP, TRUNCATE, ALTER).

Output format:
- Return the result as a valid JSON object with column names as keys.
- For computed fields (time differences, item counts), use descriptive keys
  (e.g., "prep_time_minutes", "item_count", "delivery_delay_minutes").
- If no record is found for the given order_id, return exactly: {{"error": "Order not found"}}
- Represent NULL column values as null in the JSON.
"""

sql_toolkit = SQLDatabaseToolkit(db=db, llm=model_gpt5_mini)
sql_tools = sql_toolkit.get_tools()

sql_agent = create_agent(
    model=model_gpt5_mini,
    tools=sql_tools,
    system_prompt=SQL_AGENT_SYSTEM_PROMPT,
)


def run_sql_agent_query(question: str) -> str:
    """Translate a natural language question into SQL, execute it, and return the result.

    Args:
        question (str): Natural language query about an order.

    Returns:
        str: JSON string with order data, or an error message if not found.
    """
    events = list(
        sql_agent.stream(
            {"messages": [{"role": "user", "content": question}]},
            stream_mode="values",
        )
    )
    final_message = events[-1]["messages"][-1]
    content = final_message.content
    # Reasoning models return content as a list of blocks; extract text only.
    if isinstance(content, list):
        return " ".join(block["text"] for block in content if block.get("type") == "text")
    return content


# ---------------------------------------------------------------------------
# Answer Tool Prompt — output-side guardrails
# ---------------------------------------------------------------------------

ANSWER_TOOL_PROMPT = """
You are FoodHub's customer response specialist. Your role is to transform raw \
order data and the customer's original question into a clear, polite, and \
professional reply.

Strict Focus Rule:
- Answer ONLY the specific question asked. Nothing more.
- Do NOT include any field or detail that was not directly asked for.
- Do NOT add bullet points, summaries, or extra context unless the customer \
explicitly asked for them.
- Do NOT make proactive offers, suggestions, or ask follow-up questions.
- Example: if the customer asked "what is the status?", reply with the status \
in one sentence and stop.

Tone and Format:
- Warm, empathetic, and professional — the customer may be anxious or frustrated
- Plain English — do not expose internal field names (e.g. say "your order" \
not "order_id") or raw JSON to the customer

Accuracy Guardrails:
- Only reference information explicitly present in the raw order data provided
- Do not invent, estimate, or infer values that are missing or null
- If delivery_time is null or missing, state that the order has not yet been \
delivered — do not guess an ETA beyond what the data shows
- If the raw data indicates an error or order not found, apologize politely and \
ask the customer to verify their Order ID

Policy Guardrails:
- Do not make policy promises (e.g. "you will receive a refund within X hours") \
unless that information is explicitly present in the raw data
- Do not include external delivery carrier tracking links

Escalation:
- If the raw data or customer query context indicates an unresolved repeated \
complaint, account access issue, or address change request, include this exact \
phrase in your response:
  "I'm escalating your query to a human agent who will assist you shortly."
"""

# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@tool(
    "order_query_tool",
    description=(
        "Query the FoodHub order database for order-related information. "
        "Use this tool when the customer asks about order status, delivery ETA, "
        "items in the order, payment status, preparation time, or any other "
        "order-specific details. Include the Order ID (e.g. 'O12486') in the "
        "query when known for a precise database lookup."
    ),
)
def order_query_tool(
    query: Annotated[
        str,
        (
            "The customer's natural language question about their order. "
            "Should include the Order ID (format: letter + digits, e.g. 'O12486') "
            "when known for a precise lookup. "
            "Examples: 'What is the status of order O12486?', "
            "'Is order O99123 delivered?', 'What items are in order O55312?'"
        ),
    ],
) -> str:
    """Query the FoodHub order database and return raw order data."""
    return run_sql_agent_query(query)


@tool(
    "answer_tool",
    description=(
        "Refine raw order data into a polite, formal, customer-ready response. "
        "Call this tool after order_query_tool has returned raw order data. "
        "Requires the customer's original question and the raw order data — "
        "both are needed to produce a relevant, data-grounded reply."
    ),
)
def answer_tool(
    customer_query: Annotated[
        str,
        (
            "The customer's original question, exactly as submitted. "
            "Used to ensure the response directly addresses what was asked."
        ),
    ],
    raw_order_data: Annotated[
        str,
        (
            "The raw order data string returned by order_query_tool. "
            "Should be the complete JSON response from the database."
        ),
    ],
) -> str:
    """Compose a polite, customer-ready response from raw order data."""
    messages = [
        SystemMessage(content=ANSWER_TOOL_PROMPT),
        HumanMessage(
            content=(
                f"Customer question: {customer_query}\n\n"
                f"Raw order data from database:\n{raw_order_data}"
            )
        ),
    ]
    response = model_gpt5_mini.invoke(messages)
    return response.content


# ---------------------------------------------------------------------------
# Chat Agent
# ---------------------------------------------------------------------------

chat_tools = [order_query_tool, answer_tool]

CHAT_AGENT_SYSTEM_PROMPT = """
You are FoodHub's AI customer service assistant. You help customers with \
queries about their food orders by fetching accurate information from the \
order database and providing polite, professional responses.

Tool Usage — Always follow this sequence for order-related queries:
1. Call order_query_tool with the customer's question to retrieve raw order data
2. Call answer_tool with the customer's original question AND the raw data \
returned by order_query_tool to compose the final customer reply
3. Never answer order-specific questions (status, ETA, items, payment) from \
memory — always fetch from the database first

Input Guardrails:

Security and Scope:
- If a user claims to be a hacker, requests data for all orders, or attempts \
any form of unauthorized access, refuse politely and do not call any tools
- Only assist with FoodHub order-related queries (order status, delivery ETA, \
items, payment, cancellations). For unrelated topics, politely decline and \
redirect the customer to the appropriate channel

Order ID Requirement:
- If the customer asks about a specific order but has not provided an Order ID, \
ask for it before calling order_query_tool
- Order IDs follow the format: a letter followed by digits (e.g. O12486)

Escalation:
- If the customer states their issue has not been resolved after multiple \
attempts, escalate immediately to a human agent
- If the customer requests an address change or account-related modification, \
escalate to a human agent — do not attempt to make changes
- When escalating, use this exact phrase: \
"I'm escalating your query to a human agent who will assist you shortly."
"""

chat_agent = create_agent(
    model=model_gpt5_mini,
    tools=chat_tools,
    system_prompt=CHAT_AGENT_SYSTEM_PROMPT,
)


def run_chat_agent_query(query: str) -> str:
    """Pass a customer query to the Chat Agent and return the final response.

    Args:
        query (str): The customer's natural language question or request.

    Returns:
        str: The final customer-facing response from the Chat Agent.
    """
    events = list(
        chat_agent.stream(
            {"messages": [{"role": "user", "content": query}]},
            stream_mode="values",
            config={"recursion_limit": 50},
        )
    )
    content = events[-1]["messages"][-1].content
    if isinstance(content, list):
        return " ".join(
            block["text"] for block in content if block.get("type") == "text"
        )
    return content