import os, json
import chromadb
import gradio as gr
from sentence_transformers import SentenceTransformer
from groq import Groq

# ── Chargement modèle embeddings ─────────────────────────────
print("⏳ Loading embeddings model...")
embedder = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
print("✅ Embeddings ready")

# ── Chargement chunks.json → ChromaDB en mémoire ─────────────
print("⏳ Loading chunks.json...")
with open("chunks.json", "r", encoding="utf-8") as f:
    chunks = json.load(f)

chroma_client = chromadb.Client()
collection = chroma_client.create_collection("startupfinance_tn")

embeddings_list = embedder.encode([c["content"] for c in chunks]).tolist()
collection.add(
    ids=[str(i) for i in range(len(chunks))],
    embeddings=embeddings_list,
    documents=[c["content"] for c in chunks],
    metadatas=[{"titre": c["titre"], "domaine": c["domaine"], "source": c["source"]} for c in chunks]
)
print(f"✅ ChromaDB ready — {len(chunks)} chunks indexed")

# ── Client Groq ───────────────────────────────────────────────
groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

SYSTEM_PROMPT = """You are FiskoBot, an expert financial assistant specialized in the Tunisian startup ecosystem.
You help founders, entrepreneurs and investors understand:
- The financial benefits of the Tunisian Startup Act
- The fiscal provisions of the Finance Law 2026
- Financing mechanisms, tax exemptions, and available support

Rules:
1. Answer ONLY based on the provided context.
2. If the information is missing, say so clearly.
3. Always cite the source.
4. Be precise about figures and conditions.
5. Always reply in the same language as the user's question. If the user writes in French, reply in French. If in English, reply in English. If in Arabic, reply in Arabic."""

def ask_rag(question: str, history: list = []):
    q_embedding = embedder.encode([question]).tolist()
    results = collection.query(query_embeddings=q_embedding, n_results=5)
    docs = results["documents"][0]
    metas = results["metadatas"][0]
    context = "\n\n---\n\n".join([f"[{metas[i]['titre']}]\n{docs[i]}" for i in range(len(docs))])
    sources = list({m["titre"] for m in metas})

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        if h["role"] in ["user", "assistant"]:
            messages.append({"role": h["role"], "content": h["content"]})
    lang = "French" if any(c in question for c in "àâäéèêëîïôùûüç") else "English"
    messages.append({"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}\n\nIMPORTANT: You MUST reply in {lang} only."})

    resp = groq_client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=messages,
        max_tokens=1500,
        temperature=0.1
    )
    answer = resp.choices[0].message.content
    if sources:
        answer += "\n\n---\n📚 **Sources:** " + " · ".join(sources)
    return answer

# ── Gradio Interface ──────────────────────────────────────────
SUGGESTIONS = [
    "What are the tax benefits for a labeled startup?",
    "How to obtain the Startup Act label in Tunisia?",
    "What is the 100k TND technology card?",
    "What benefits do investors in startups get?",
    "What are the SME measures in Finance Law 2026?",
]

def respond(message, history):
    if not message.strip():
        return "", history
    answer = ask_rag(message, history)
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": answer})
    return "", history

INITIAL = [{"role": "assistant", "content": "👋 Hello! I'm **FiskoBot**.\nAsk me anything about the **Startup Act** or the **Finance Law 2026**!"}]

with gr.Blocks(title="FiskoBot 🇹🇳") as demo:
    gr.HTML("""
    <div style="text-align:center;padding:20px 0 10px;background:linear-gradient(135deg,#1e1b4b,#4c1d95);border-radius:12px;margin-bottom:16px">
        <h1 style="font-size:2rem;margin:0;color:white">🇹🇳 FiskoBot</h1>
        <p style="color:#c4b5fd;margin:6px 0 0">AI Financial Assistant — Startup Act & Finance Law 2026</p>
    </div>""")

    chatbot = gr.Chatbot(height=450, value=INITIAL)

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask your question about the Startup Act, taxation...",
            show_label=False,
            scale=5,
            container=False
        )
        btn = gr.Button("Send ▶", variant="primary", scale=1)

    gr.Button("🗑️ New conversation", size="sm").click(
        lambda: (INITIAL, ""),
        outputs=[chatbot, msg]
    )

    gr.Markdown("### 💡 Frequently asked questions")
    with gr.Row():
        for q in SUGGESTIONS[:3]:
            gr.Button(q[:50]+"..." if len(q)>50 else q, size="sm").click(fn=lambda x=q: x, outputs=msg)
    with gr.Row():
        for q in SUGGESTIONS[3:]:
            gr.Button(q[:50]+"..." if len(q)>50 else q, size="sm").click(fn=lambda x=q: x, outputs=msg)

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    btn.click(respond, [msg, chatbot], [msg, chatbot])

demo.launch(server_name="0.0.0.0", server_port=7860)