import streamlit as st
import sqlite3
import pandas as pd
import json
import re
import os
from datetime import date
from typing import TypedDict, List, Dict, Any

from openai import OpenAI
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
from langchain_core.tools import tool

# ── Page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Kartify Support Hub",
    page_icon="🛒",
    layout="centered",
)

# ── LLMs ─────────────────────────────────────────────────────────────────────
@st.cache_resource
def load_llms():
    llm          = ChatOpenAI(model_name="gpt-4o-mini")
    evaluate_llm = ChatOpenAI(model_name="gpt-4o")
    return llm, evaluate_llm

llm, evaluate_llm = load_llms()

# ── State ─────────────────────────────────────────────────────────────────────
class OrderState(TypedDict):
    cust_id:          str
    order_id:         str
    order_context:    str
    query:            str
    raw_agent_response: str
    final_response:   str
    history:          List[Dict[str, str]]
    intent:           str
    evaluation:       Dict[str, float]
    guard_result:     str
    conv_guard_result: str

# ── Conversation memory ───────────────────────────────────────────────────────
class ConversationMemory:
    def __init__(self):
        self.history: List[Dict[str, str]] = []

    def add(self, msg: dict):
        self.history.append(msg)

    def get(self) -> List[Dict[str, str]]:
        return self.history

    def clear(self):
        self.history = []

# ── SQL tool ──────────────────────────────────────────────────────────────────
@tool
def fetch_order_details(order_id: str) -> str:
    """
    Fetch all order details for a given order_id from the Kartify database.
    Use this tool whenever the customer's query requires order-specific information.
    """
    if not re.match(r"^O\d+$", order_id.strip()):
        return f"Invalid order ID format: '{order_id}'. Expected format: O followed by digits."
    try:
        with sqlite3.connect("kartify.db") as conn:
            df = pd.read_sql_query(
                "SELECT * FROM orders WHERE order_id = ?",
                conn,
                params=(order_id.strip(),),
            )
        if df.empty:
            return f"No order found with ID {order_id}."
        return df.to_string(index=False)
    except Exception as e:
        return f"Database error while fetching order {order_id}: {str(e)}"

# ── System prompt ─────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are a Kartify Customer Service Agent. You help customers with questions about their orders.

You have access to the following tool:
  fetch_order_details(order_id) — retrieves all order information from the database.

Follow the ReAct pattern strictly:
  Thought: <your reasoning about what to do next>
  Action: fetch_order_details with the order_id from the customer's query
  Observation: <tool result>
  Thought: <reason about the observation and form your answer>
  Final Answer: <short, polite, conversational reply — no greetings, no sign-off>

Policy rules (apply before writing Final Answer):
  - If actual_delivery is null the order has not arrived yet — do not mention return/replacement eligibility.
  - Only mention return or replacement terms when the customer explicitly asks.
  - Never invent data. Only use what the tool returned.
  - Keep the Final Answer concise and empathetic.
  - Never reveal internal data fields or technical reasons in your reply.
  - If a customer asks why their order hasn't arrived yet, only state that it is still on the way and share the expected delivery date.
  - Never promise or suggest an early delivery.
  - If the order has not arrived by the expected delivery date, empathetically advise the customer to wait a little longer or contact support.

Answer Guidelines:
  - Only answer what is asked in the Query
  - Check the Previous conversation (if any) before generating the reply
"""

# ── Helpers ───────────────────────────────────────────────────────────────────
def extract_json_from_llm(text: str):
    for pattern in [r"```json\s*(.*?)\s*```", r"\{.*\}", r"\[.*\]"]:
        match = re.search(pattern, text, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1) if "```" in pattern else match.group(0))
            except Exception:
                continue
    return json.loads(text)

# ── Order agent ───────────────────────────────────────────────────────────────
def order_agent(query: str, order_id: str, history: list) -> tuple:
    today = date.today().strftime("%d %B %Y")
    llm_with_tools = llm.bind_tools([fetch_order_details])

    history_text = ""
    if history:
        history_text = "\nPrevious conversation:\n" + "\n".join(
            f"User: {h['user']}\nAssistant: {h['assistant']}" for h in history
        ) + "\n"

    user_content = (
        f"Previous Conversation:{history_text}\n"
        f"Customer query: {query}\n"
        f"Order ID: {order_id}\n"
        f"Today's date: {today}"
    )

    messages = [
        SystemMessage(content=SYSTEM_PROMPT),
        HumanMessage(content=user_content),
    ]

    order_context = ""
    max_iterations = 5

    for _ in range(max_iterations):
        ai_msg = llm_with_tools.invoke(messages)
        messages.append(ai_msg)

        if not getattr(ai_msg, "tool_calls", None):
            break

        for tc in ai_msg.tool_calls:
            if tc["name"] == "fetch_order_details":
                result = fetch_order_details.invoke(tc["args"])
                order_context = result
                messages.append(ToolMessage(content=result, tool_call_id=tc["id"]))

    final_response = ai_msg.content.strip()
    for prefix in ("Final Answer:", "final answer:"):
        if final_response.lower().startswith(prefix.lower()):
            final_response = final_response[len(prefix):].strip()
            break

    return order_context, final_response

# ── Node functions ────────────────────────────────────────────────────────────
def user_input_node(state: OrderState): return state
def memory_node(state: OrderState):
    st.session_state.conversation_memory.add({"user": state["query"], "assistant": state["final_response"]})
    return state
def order_agent_node(state: OrderState):
    oc, fr = order_agent(query=state["query"], order_id=state["order_id"], history=state["history"])
    return {"order_context": oc, "final_response": fr}
def intent_node(state: OrderState):
    prompt = f"Classify intent into numeric ID (0, 1, 2, 3) only:\nQuery: {state['query']}"
    result = llm.invoke([HumanMessage(content=prompt)]).content.strip()
    return {"intent": result[:1]}
def router_node(state: OrderState): return "order_agent" if state["intent"] == "2" else "exit_node"
def exit_node(state: OrderState):
    mapping = {
        "0": "Sorry for the inconvenience. A human support agent will assist you shortly.",
        "1": "Thank you! I hope I was able to assist with your query.",
        "3": "Apologies, I'm currently only able to help with information about your placed orders.",
    }
    return {"final_response": mapping.get(state["intent"], "How can I help you?")}
def evaluation_node(state: OrderState):
    prompt = f"Evaluate response JSON format only:\nContext: {state['order_context']}\nQuery: {state['query']}\nResponse: {state['final_response']}"
    try:
        raw = evaluate_llm.invoke([HumanMessage(content=prompt)]).content.strip()
        evaluation = extract_json_from_llm(raw)
    except Exception: evaluation = {"groundedness": 1.0, "precision": 1.0}
    return {"evaluation": evaluation}
def retry_router(state: OrderState):
    score = state.get("evaluation", {})
    if score.get("groundedness", 0) < 0.75 or score.get("precision", 0) < 0.75: return "order_agent"
    return "safety_check"
def guard_node(state: OrderState):
    prompt = f"Classify content BLOCK or SAFE:\nResponse: {state['final_response']}"
    result = evaluate_llm.invoke([HumanMessage(content=prompt)]).content.strip()
    guard_result = result if result in ("BLOCK", "SAFE") else "SAFE"
    if guard_result == "BLOCK":
        return {"guard_result": guard_result, "final_response": "Your request is being forwarded to a customer support specialist."}
    return {"guard_result": guard_result}
def guard_router(state: OrderState): return "exit" if state.get("guard_result") == "BLOCK" else "memory_save"
def conversational_guard_node(state: OrderState):
    prompt = f"Review conversation safety BLOCK or SAFE:\n{state.get('history', [])}"
    result = evaluate_llm.invoke([HumanMessage(content=prompt)]).content.strip()
    conv_result = result if result in ("BLOCK", "SAFE") else "SAFE"
    if conv_result == "BLOCK":
        return {"conv_guard_result": conv_result, "final_response": "Your request is being forwarded to a customer support specialist."}
    return {"conv_guard_result": conv_result}
def conv_guard_router(state: OrderState): return "exit" if state.get("conv_guard_result") == "BLOCK" else "done"

# ── Build LangGraph ───────────────────────────────────────────────────────────
@st.cache_resource
def build_graph():
    g = StateGraph(OrderState)
    g.add_node("user_input", user_input_node)
    g.add_node("intent_classifier", intent_node)
    g.add_node("order_agent", order_agent_node)
    g.add_node("evaluate", evaluation_node)
    g.add_node("safety_check", guard_node)
    g.add_node("conv_safety_check", conversational_guard_node)
    g.add_node("memory_save", memory_node)
    g.add_node("exit_node", exit_node)
    g.set_entry_point("user_input")
    g.add_edge("user_input", "intent_classifier")
    g.add_conditional_edges("intent_classifier", router_node, {"order_agent": "order_agent", "exit_node": "exit_node"})
    g.add_edge("order_agent", "evaluate")
    g.add_conditional_edges("evaluate", retry_router, {"order_agent": "order_agent", "safety_check": "safety_check"})
    g.add_conditional_edges("safety_check", guard_router, {"memory_save": "memory_save", "exit": "exit_node"})
    g.add_edge("memory_save", "conv_safety_check")
    g.add_conditional_edges("conv_safety_check", conv_guard_router, {"done": END, "exit": "exit_node"})
    g.add_edge("exit_node", END)
    return g.compile()

order_graph = build_graph()

# ── Session state defaults ────────────────────────────────────────────────────
if "conversation_memory" not in st.session_state: st.session_state.conversation_memory = ConversationMemory()
if "chat_messages" not in st.session_state: st.session_state.chat_messages = []
if "chat_active" not in st.session_state: st.session_state.chat_active = False
if "cust_id" not in st.session_state: st.session_state.cust_id = ""
if "order_id" not in st.session_state: st.session_state.order_id = ""
if "orders_df" not in st.session_state: st.session_state.orders_df = None

def fetch_customer_orders(cust_id: str) -> pd.DataFrame | None:
    try:
        with sqlite3.connect("kartify.db") as conn:
            df = pd.read_sql_query("SELECT order_id, product_description, order_status FROM orders WHERE customer_id = ?", conn, params=(cust_id.strip(),))
        return df if not df.empty else None
    except Exception: return None

def run_turn(query: str, cust_id: str, order_id: str) -> str:
    state: OrderState = {
        "cust_id": cust_id, "order_id": order_id, "order_context": "", "query": query, "raw_agent_response": "",
        "final_response": "", "history": st.session_state.conversation_memory.get(), "intent": "", "evaluation": {},
        "guard_result": "", "conv_guard_result": "",
    }
    result = order_graph.invoke(state, config={"recursion_limit": 100})
    return result.get("final_response", "I'm sorry, I couldn't process that request.")

# ══════════════════════════════════════════════════════════════════════════════
# BRANDED UI STYLING (KARTIFY NORDIC-MODERN THEME)
# ══════════════════════════════════════════════════════════════════════════════

st.markdown(
    """
    <style>
    @import url('[https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap](https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap)');
    
    /* Global Overrides */
    html, body, [class*="css"] {
        font-family: 'Inter', sans-serif;
    }
    .block-container { max-width: 760px; padding-top: 2rem; }
    
    /* Header Branding Banner */
    .brand-banner {
        background: linear-gradient(135deg, #1E293B 0%, #0F172A 100%);
        padding: 24px;
        border-radius: 16px;
        color: white;
        margin-bottom: 25px;
        box-shadow: 0 4px 20px rgba(15, 23, 42, 0.08);
        display: flex;
        align-items: center;
        gap: 16px;
    }
    .brand-logo {
        font-size: 2.2rem;
        background: rgba(255, 255, 255, 0.1);
        padding: 8px 14px;
        border-radius: 12px;
    }
    .brand-title-text h1 {
        color: white !important;
        margin: 0 !important;
        font-size: 1.6rem !important;
        font-weight: 700;
        letter-spacing: -0.03em;
    }
    .brand-title-text p {
        color: #94A3B8 !important;
        margin: 4px 0 0 0 !important;
        font-size: 0.85rem;
    }
    
    /* Stepper UI Progress Tracker */
    .stepper-container {
        display: flex;
        justify-content: space-between;
        background: #F8FAFC;
        padding: 16px;
        border-radius: 12px;
        border: 1px solid #E2E8F0;
        margin-top: 12px;
    }
    .step-item {
        text-align: center;
        flex: 1;
        position: relative;
    }
    .step-dot {
        width: 12px;
        height: 12px;
        border-radius: 50%;
        margin: 0 auto 6px auto;
    }
    .step-dot.active { background-color: #0EA5E9; box-shadow: 0 0 0 4px rgba(14, 165, 233, 0.2); }
    .step-dot.inactive { background-color: #CBD5E1; }
    .step-label { font-size: 0.75rem; font-weight: 500; color: #64748B; }
    .step-label.active { color: #0EA5E9; font-weight: 600; }
    
    /* Sidebar Layout Polish */
    .sidebar-meta-box {
        background: #F8FAFC;
        border: 1px solid #E2E8F0;
        padding: 14px;
        border-radius: 10px;
        margin-bottom: 12px;
    }
    .sidebar-meta-label { font-size: 0.75rem; color: #64748B; text-transform: uppercase; font-weight: 600; }
    .sidebar-meta-val { font-size: 0.95rem; color: #0F172A; font-weight: 600; margin-bottom: 6px; }
    </style>
    """,
    unsafe_allow_html=True,
)

# ── Branded Header ───────────────────────────────────────────────────────────
st.markdown(
    """
    <div class="brand-banner">
        <div class="brand-logo">🛒</div>
        <div class="brand-title-text">
            <h1>KARTIFY</h1>
            <p>Premium Concierge Client Support</p>
        </div>
    </div>
    """,
    unsafe_allow_html=True
)

# ── Phase 1: Customer ID lookup ───────────────────────────────────────────────
if not st.session_state.chat_active:
    st.markdown("#### 🔑 Identity Verification")
    st.caption("Please authenticate using your structural Customer Identification profile number.")

    with st.form("customer_form"):
        cust_input = st.text_input(
            "Customer ID Token",
            placeholder="e.g. C1010",
            value=st.session_state.cust_id,
            label_visibility="collapsed"
        )
        submitted = st.form_submit_button("Verify Identity & Find Orders", use_container_width=True)

    if submitted and cust_input.strip():
        with st.spinner("Accessing global secure database records…"):
            df = fetch_customer_orders(cust_input.strip())
        if df is not None:
            st.session_state.cust_id  = cust_input.strip()
            st.session_state.orders_df = df
        else:
            st.error(f"No customer ledger files associated with account index **{cust_input.strip()}**.")

    # ── Phase 2: Order selection ──────────────────────────────────────────────
    if st.session_state.orders_df is not None:
        st.markdown("---")
        st.markdown("#### 📦 Active Orders Ledger")
        st.caption("Select an active order deployment to link your real-time conversational agent pipeline.")

        df = st.session_state.orders_df
        options = {
            f"ID: {row['order_id']} | {row['product_description'][:40]}...": row["order_id"]
            for _, row in df.iterrows()
        }

        selected_label = st.selectbox("Your orders", list(options.keys()), index=0, label_visibility="collapsed")
        selected_order_id = options[selected_label]

        # Order Micro-card Contextual Summary
        selected_row = df[df["order_id"] == selected_order_id].iloc[0]
        status = str(selected_row['order_status']).strip().lower()
        
        # Build beautiful multi-stage tracking visualizations dynamically
        st.markdown(
            f"""
            <div style="background:#FFF; border:1px solid #E2E8F0; border-radius:12px; padding:16px; margin:12px 0; box-shadow:0 1px 3px rgba(0,0,0,0.02);">
                <div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:10px;">
                    <span style="font-size:0.85rem; font-weight:700; background:#E0F2FE; color:#0369A1; padding:3px 8px; border-radius:6px;">{selected_row['order_id']}</span>
                    <span style="font-size:0.8rem; color:#64748B;">Item: <strong>{selected_row['product_description']}</strong></span>
                </div>
                <div class="stepper-container">
                    <div class="step-item">
                        <div class="step-dot {'active' if status in ['ordered', 'processing', 'shipped', 'delivered'] else 'inactive'}"></div>
                        <div class="step-label {'active' if status=='ordered' else ''}">Ordered</div>
                    </div>
                    <div class="step-item">
                        <div class="step-dot {'active' if status in ['processing', 'shipped', 'delivered'] else 'inactive'}"></div>
                        <div class="step-label {'active' if status=='processing' else ''}">Processing</div>
                    </div>
                    <div class="step-item">
                        <div class="step-dot {'active' if status in ['shipped', 'delivered'] else 'inactive'}"></div>
                        <div class="step-label {'active' if status=='shipped' else ''}">In Transit</div>
                    </div>
                    <div class="step-item">
                        <div class="step-dot {'active' if status=='delivered' else 'inactive'}"></div>
                        <div class="step-label {'active' if status=='delivered' else ''}">Delivered</div>
                    </div>
                </div>
            </div>
            """,
            unsafe_allow_html=True,
        )

        if st.button("Initialize Secure Chat Channel", use_container_width=True, type="primary"):
            st.session_state.order_id = selected_order_id
            st.session_state.chat_active = True
            st.session_state.conversation_memory.clear()
            st.session_state.chat_messages = []
            st.session_state.chat_messages.append({
                "role": "assistant",
                "content": f"Welcome to Kartify Concierge Service. I have securely retrieved parameters for order **{selected_order_id}**. What context or tracking diagnostics can I deliver for you today?"
            })
            st.rerun()

# ── Phase 3: Chat interface ───────────────────────────────────────────────────
else:
    # Sidebar Session Details
    with st.sidebar:
        st.markdown("### 🛒 Session Parameters")
        st.markdown(
            f"""
            <div class="sidebar-meta-box">
                <div class="sidebar-meta-label">Client Token</div>
                <div class="sidebar-meta-val">{st.session_state.cust_id}</div>
                <div class="sidebar-meta-label">Linked Context</div>
                <div class="sidebar-meta-val">{st.session_state.order_id}</div>
            </div>
            """,
            unsafe_allow_html=True
        )
        
        if st.button("Disconnect Session", use_container_width=True):
            st.session_state.chat_active = False
            st.session_state.chat_messages = []
            st.session_state.conversation_memory.clear()
            st.session_state.orders_df = None
            st.session_state.cust_id = ""
            st.session_state.order_id = ""
            st.rerun()
            
        st.divider()
        st.caption("Kartify Conversational Framework v2.1\n\nSecurity Model: Active Guardrails Enabled.")

    # Render Active Messaging Stream
    for msg in st.session_state.chat_messages:
        if msg["role"] == "user":
            with st.chat_message("user"):
                st.markdown(msg["content"])
        else:
            with st.chat_message("assistant", avatar="🛒"):
                st.markdown(msg["content"])

    # Chat User Prompt Field
    user_query = st.chat_input("Inquire about shipping data, tracking parameters, or item diagnostics...")

    if user_query:
        st.session_state.chat_messages.append({"role": "user", "content": user_query})
        with st.chat_message("user"):
            st.markdown(user_query)

        with st.chat_message("assistant", avatar="🛒"):
            with st.spinner("Analyzing data engine graphs..."):
                response = run_turn(
                    query=user_query,
                    cust_id=st.session_state.cust_id,
                    order_id=st.session_state.order_id,
                )
            st.markdown(response)

        st.session_state.chat_messages.append({"role": "assistant", "content": response})

        # Automatic termination visual indicator rules
        exit_phrases = ["human support agent", "customer support specialist", "I hope I was able to assist", "only able to help with information"]
        if any(p.lower() in response.lower() for p in exit_phrases):
            st.info("System Notification: Session routing successfully complete. Active pipeline locked.")