import os
import re
import gradio as gr

from langchain_groq import ChatGroq
from langchain_core.messages import SystemMessage

# ✅ Correct imports for SQLDatabase + Toolkit (HF-safe)
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit

# ✅ Correct import path for create_sql_agent (NOT langchain.agents)
from langchain_community.agent_toolkits.sql.base import create_sql_agent


# -----------------------------
# Config
# -----------------------------
MODEL_NAME = os.getenv("GROQ_MODEL", "meta-llama/llama-3.1-8b-instant")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError(
        "❌ GROQ_API_KEY not found. Add it in Hugging Face: Settings → Secrets → GROQ_API_KEY"
    )

DB_PATH = "customer_orders.db"  # must be in same Space repo
DB_URI = f"sqlite:///{DB_PATH}"

if not os.path.exists(DB_PATH):
    raise FileNotFoundError(
        f"❌ Database file '{DB_PATH}' not found. Make sure customer_orders.db is committed to the Space repo."
    )

# -----------------------------
# LLM
# -----------------------------
llm = ChatGroq(
    model=MODEL_NAME,
    temperature=0,      # deterministic
    max_tokens=300,
    groq_api_key=GROQ_API_KEY
)

# -----------------------------
# Database + SQL Agent
# -----------------------------
db = SQLDatabase.from_uri(DB_URI)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)

system_message = SystemMessage(content=
    "You are FoodHub customer support. "
    "You can query the database safely using SQL tools. "
    "Rules: (1) Never return bulk data. (2) Only answer for a single order_id. "
    "(3) If order_id missing, ask for it. (4) Keep replies short and professional."
)

sql_agent = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=False,  # set True only if you want logs
    system_message=system_message
)

ORDER_ID_PATTERN = re.compile(r"\bO\d{5}\b", re.IGNORECASE)

def extract_order_id(text: str):
    m = ORDER_ID_PATTERN.search(text or "")
    return m.group(0).upper() if m else None


# -----------------------------
# Chatbot logic
# -----------------------------
def respond(user_message, history):
    user_message = (user_message or "").strip()

    # 1) If order id not provided, ask for it
    order_id = extract_order_id(user_message)
    if not order_id:
        return "Please provide your Order ID (example: O12488)."

    # 2) Ask SQL agent for that one order only (safe query)
    # Keep query explicit to avoid the agent doing broad selects
    agent_query = (
        f"Retrieve all columns for order_id '{order_id}' from the orders table. "
        f"Then summarize the order status and payment status in 2-3 short sentences."
    )

    try:
        result = sql_agent.invoke({"input": agent_query})
        # result is usually dict with "output"
        output_text = result.get("output", str(result))
        return output_text
    except Exception as e:
        return f"Sorry — I couldn’t fetch that order right now. Please try again. (Error: {type(e).__name__})"


# -----------------------------
# Gradio UI
# -----------------------------
demo = gr.ChatInterface(
    fn=respond,
    title="FoodHub – AI Powered Food Delivery Chatbot",
    description="Ask order-related questions. Include your Order ID like O12488.",
    examples=[
        "Where is my order? O12488",
        "What is the payment status of order O12488?",
        "I want to cancel my order O12487"
    ],
)

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