import streamlit as st
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer
import requests
import sqlite3

# 🧠 Load embedding model
@st.cache_resource
def load_model():
    return SentenceTransformer("all-MiniLM-L6-v2")

model = load_model()

# 🔗 Connect to Qdrant (FAQ vector DB)
client = QdrantClient(host="x", port=6333)

# 🔗 Connect to your live database (SQLite example)
conn = sqlite3.connect("logistics.db")

# 🤖 Helper to call Mistral API
def query_mistral(prompt):
    try:
        res = requests.post("http://localhost:11434/api/generate", json={
            "model": "mistral",
            "prompt": prompt,
            "stream": False
        })
        return res.json().get("response", "").strip()
    except Exception as e:
        return f"⚠️ Error connecting to Mistral: {e}"

# 🎯 Step 1: Classify user input intent with improved prompt
def classify_input(user_input):
    prompt = (
        "You are a classification model.\n"
        "Classify the following user input into exactly one of two categories:\n"
        " - FAQ : if the user is asking general questions about logistics or ShypBuddy service.\n"
        " - LIVE_DATA : if the user is asking about their specific shipment status, tracking info, or live data retrieval.\n"
        "\n"
        "Only respond with exactly one word: FAQ or LIVE_DATA.\n"
        "\n"
        f"User Input: \"{user_input}\"\n"
        "Classification:"
    )
    classification = query_mistral(prompt).strip().upper()
    if classification not in ["FAQ", "LIVE_DATA"]:
        classification = "UNKNOWN"
    return classification

# 🔍 FAQ semantic search
def search_faq(query):
    query_vector = model.encode([query])[0].tolist()
    results = client.search(
        collection_name="shypbuddy_faq_embeddings",
        query_vector=query_vector,
        limit=3,
        with_payload=True,
        score_threshold=0.7
    )
    return results

# 🔄 Step 2: Generate SQL query from user input using Mistral
def generate_sql_query(user_input):
    prompt = (
        "You are an assistant that converts user shipment tracking or logistics queries into safe SQL queries. "
        "Only generate SELECT queries on the shipments table with columns: tracking_number, status, location, eta.\n"
        "If you cannot generate SQL, reply with 'NO_SQL'.\n\n"
        f"User question: \"{user_input}\"\n"
        "SQL query:"
    )
    sql_query = query_mistral(prompt)
    if "NO_SQL" in sql_query or not sql_query.strip().upper().startswith("SELECT"):
        return None
    return sql_query.strip()

# 🗄 Execute SQL and format response
def query_live_data(user_input):
    sql_query = generate_sql_query(user_input)
    if not sql_query:
        return (
            "❗ Sorry, I couldn't generate a valid shipment tracking SQL query. "
            "Please specify your tracking number."
        )

    try:
        cursor = conn.cursor()
        cursor.execute(sql_query)
        row = cursor.fetchone()
        if not row:
            return "❗ No shipment found matching your query."
        # Adjust according to your actual SELECT column order
        tracking_number, status, location, eta = row
        response = (
            f"Shipment {tracking_number} is currently '{status}', "
            f"located at '{location}', with an estimated arrival time of {eta}."
        )
        return response
    except Exception as e:
        return f"⚠️ Error running SQL query: {e}"

# --- Streamlit UI ---
st.title("📦 ShypBuddy Logistic Chatbot (FAQ + Live Data with Mistral)")

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

with st.form("chat_form", clear_on_submit=True):
    user_input = st.text_input("Enter your logistics question or shipment query", key="user_input")
    submitted = st.form_submit_button("Send")

if submitted and user_input:
    st.session_state.chat_history.append(("user", user_input))
    classification = classify_input(user_input)

    # DEBUG: show classification result in UI (can remove or comment out after debugging)
    st.write(f"**DEBUG: Classification result = '{classification}'**")

    if classification == "LIVE_DATA":
        reply = query_live_data(user_input)
        st.session_state.chat_history.append(("bot", reply))

    elif classification == "FAQ":
        matches = search_faq(user_input)

        if matches:
            top_score = matches[0].score
            top_payload = matches[0].payload

            if top_score >= 0.85:
                # High confidence, answer directly
                st.session_state.chat_history.append(("bot", f"✅ {top_payload['answer']}"))
            else:
                # Vague question, show list of FAQs to choose from
                st.session_state.chat_history.append(("bot", "🔍 I found some related FAQs, please select one:"))
                st.session_state.last_matches = matches
        else:
            st.session_state.chat_history.append(("bot", "🚫 Sorry, I couldn't find any related logistics FAQs."))

    else:
        st.session_state.chat_history.append(("bot", "🚫 Sorry, I can only answer logistics FAQs or shipment tracking questions."))

# Display chat history
for role, message in st.session_state.chat_history:
    if role == "user":
        st.markdown(f"🧑‍💻 **You**: {message}")
    else:
        st.markdown(f"🤖 **Bot**: {message}")

# Show FAQ choices if available
if "last_matches" in st.session_state:
    for idx, match in enumerate(st.session_state.last_matches, 1):
        if st.button(f"{idx}. {match.payload['question']}"):
            st.session_state.chat_history.append(("user", match.payload['question']))
            st.session_state.chat_history.append(("bot", f"✅ Answer: {match.payload['answer']}"))
            del st.session_state.last_matches
            st.experimental_rerun()

For hi or hello kind of question return something meaning full

if vague FAQ is asked suggest the related FAQ