import os
import traceback
from pathlib import Path

import gradio as gr

from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit


DB_PATH = Path("database_") / "zain_customer_360_ai_demo.db"
MODEL_NAME = "gpt-4.1-mini"


SYSTEM_PROMPT = """
You are a telecom business intelligence SQL agent for Zain Jordan.

You are connected to a SQLite Customer 360 database.

Rules:
- Use SELECT queries only.
- Never modify the database.
- Never use INSERT, UPDATE, DELETE, DROP, ALTER, TRUNCATE, CREATE, REPLACE, ATTACH, DETACH, or VACUUM.
- Inspect the database schema before writing SQL.
- Limit results to 10 rows unless the user asks otherwise.
- Use only relevant columns.
- Do not guess facts that are not in the database.
- Explain the result in simple business language.

Useful context:
- Churn analysis usually uses customers and customer_churn_scores.
- Revenue/value analysis usually uses customer_value_segments, invoices, payments, or transactions.
- Complaint analysis usually uses complaints and support_interactions.
- Campaign analysis usually uses campaigns and customer_campaign_responses.
- Network analysis usually uses network_towers and network_events.

Final answer format:
1. Direct Answer
2. Key Numbers
3. Business Interpretation
4. Recommended Next Action
""".strip()


sql_agent = None


def get_sql_agent():
    global sql_agent

    if sql_agent is not None:
        return sql_agent

    api_key = os.environ.get("OPENAI_API_KEY")

    if not api_key:
        raise ValueError(
            "OPENAI_API_KEY is missing. Add it in Hugging Face Space Settings → Variables and secrets → Secrets."
        )

    if not DB_PATH.exists():
        raise FileNotFoundError(
            f"Database not found at: {DB_PATH}\n\n"
            "Please create a folder named database_ and upload zain_customer_360_ai_demo.db inside it."
        )

    db_uri = f"sqlite:///{DB_PATH.resolve()}"
    db = SQLDatabase.from_uri(db_uri)

    llm = init_chat_model(
        MODEL_NAME,
        model_provider="openai"
    )

    toolkit = SQLDatabaseToolkit(
        db=db,
        llm=llm
    )

    tools = toolkit.get_tools()

    sql_agent = create_agent(
        model=llm,
        tools=tools,
        system_prompt=SYSTEM_PROMPT
    )

    return sql_agent


def run_sql_agent(question):
    try:
        if not question or not question.strip():
            return "Please enter a business question."

        agent = get_sql_agent()

        result = agent.invoke({
            "messages": [
                {
                    "role": "user",
                    "content": question
                }
            ]
        })

        final_message = result["messages"][-1]

        if hasattr(final_message, "content"):
            return final_message.content

        return str(final_message)

    except Exception as e:
        return f"""
### Error

{str(e)}

### Traceback

```text
{traceback.format_exc()}
```
"""


with gr.Blocks(title="Zain Jordan SQL Agent") as demo:

    gr.Markdown("# Zain Jordan Customer 360 SQL Agent")

    gr.Markdown("""
Ask business questions from the Zain Jordan Customer 360 SQLite database.

This app uses:

- SQLite database
- LangChain `SQLDatabase`
- LangChain `SQLDatabaseToolkit`
- OpenAI `gpt-4.1-mini`
- Gradio interface

Expected database location:

```text
database_/zain_customer_360_ai_demo.db
```

Example questions:

- Which cities have the most high-risk churn customers?
- Which customer value segments have the highest average ARPU?
- What are the top complaint categories?
- Which campaigns had the highest conversion rate?
""")

    question = gr.Textbox(
        label="Ask a business question",
        lines=4,
        value="Which cities have the most high-risk churn customers? Show the top 10 cities.",
        placeholder="Ask a question about churn, revenue, complaints, campaigns, customers, or network events."
    )

    run_button = gr.Button("Run SQL Agent", variant="primary")

    answer = gr.Markdown(label="Answer")

    run_button.click(
        fn=run_sql_agent,
        inputs=question,
        outputs=answer
    )


if __name__ == "__main__":
    demo.launch()