# STARTUP:
# 1. pip install -r requirements.txt
# 2. python app.py      ← start the app (auto-builds index)

import os
import io
import numpy as np
import soundfile as sf
import gradio as gr

import stt
import llm
import tts
import rag_pipeline

# ─────────────────────────────────────────────
# AUTO INGEST — builds ChromaDB on first startup
# Runs automatically if chroma_db folder not found
# ─────────────────────────────────────────────
import os
if not os.path.exists("./chroma_db"):
    print("ChromaDB not found — building knowledge base index...")
    try:
        from knowledge_base import KNOWLEDGE_BASE
        from sentence_transformers import SentenceTransformer
        import chromadb

        print("Loading embedding model for ingest...")
        _embedder = SentenceTransformer(
            'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'
        )
        _chroma = chromadb.PersistentClient(path="./chroma_db")
        _collection = _chroma.get_or_create_collection("banking_hindi")

        _documents = []
        _metadatas = []
        _ids = []

        for doc in KNOWLEDGE_BASE:
            full_text = f"{doc['title']}\n{doc['content']}"
            _documents.append(full_text)
            _metadatas.append({
                "id": doc["id"],
                "title": doc["title"],
                "category": doc["category"]
            })
            _ids.append(doc["id"])

        print(f"Embedding {len(_documents)} documents... (takes 2-3 min on first run)")
        _embeddings = _embedder.encode(
            _documents,
            show_progress_bar=True,
            batch_size=8
        ).tolist()

        _collection.add(
            documents=_documents,
            embeddings=_embeddings,
            metadatas=_metadatas,
            ids=_ids
        )
        print(f"✅ ChromaDB ready — {len(_documents)} documents indexed")
        
        # Cleanup temp variables
        del _embedder, _chroma, _collection
        del _documents, _metadatas, _ids, _embeddings

    except Exception as e:
        print(f"⚠️ Auto-ingest failed: {e}. Will use keyword retrieval fallback.")
else:
    print("✅ ChromaDB found — skipping ingest")

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────
HF_TOKEN = (
    os.environ.get("HF_TOKEN") or
    os.environ.get("HF_API_TOKEN") or
    os.environ.get("HUGGINGFACE_TOKEN") or
    ""
)
print(f"DEBUG app.py: HF_TOKEN loaded = {bool(HF_TOKEN)}, length = {len(HF_TOKEN)}")


# (STT, LLM, and TTS functions moved to separate modules)


# ─────────────────────────────────────────────
# PIPELINE GENERATORS
# ─────────────────────────────────────────────

def run_voice_pipeline(audio_input):
    """Voice mode: mic audio → transcript, answer, citations, audio, status."""
    if audio_input is None:
        yield "", "कृपया माइक्रोफोन बटन दबाकर अपना प्रश्न पूछें।", "", None, "🔴 *तैयार*"
        return
    if not HF_TOKEN:
        yield "", "⚠️ HF_TOKEN Secret नहीं मिला। Space Settings → Secrets में जोड़ें।", "", None, "❌ *त्रुटि*"
        return

    try:
        yield "", "", "", None, "🎙️ **सुन रहा हूँ...** (STT)"
        sample_rate, audio_array = audio_input
        
        # Ensure STT gets proper data
        transcript = stt.stt_whisper(audio_array, sample_rate)
        
        if not transcript.strip():
            yield "", "आवाज़ समझ नहीं आई। कृपया फिर से बोलें।", "", None, "⚠️ *फिर से प्रयास करें*"
            return
            
        yield transcript, "", "", None, "🔎 **जानकारी खोज रहा हूँ...** (RAG)"
        normalized = rag_pipeline.normalize_jargon(transcript)
        eng_query  = rag_pipeline.translate_to_retrieval_query(normalized)
        context_docs = rag_pipeline.retrieve_with_embeddings(eng_query, top_k=3)
        
        citations = "\n".join([f"• {doc['title']}" for doc in context_docs]) if context_docs else ""
        
        yield transcript, "", "", None, "🤖 **उत्तर तैयार कर रहा हूँ...** (LLM)"
        
        context = "\n\n".join(
            f"[{i+1}] {d['title']}\n{d['content'][:500]}"
            for i, d in enumerate(context_docs)
        ) if context_docs else "कोई प्रासंगिक जानकारी नहीं मिली।"

        full_prompt = (
            f"नीचे दी गई जानकारी के आधार पर प्रश्न का उत्तर दें:\n\n"
            f"{context}\n\n"
            f"प्रश्न: {transcript}"
        )
        answer = llm.llm_generate(full_prompt)
        
        yield transcript, answer, citations, None, "🔊 **आवाज़ बना रहा हूँ...** (TTS)"
        clean_text = rag_pipeline.format_response_for_tts(answer)
        audio_out = tts.tts_hindi(clean_text)
        
        yield transcript, answer, citations, audio_out, "✅ **पूरा हुआ**"
        
    except Exception as e:
        import traceback; traceback.print_exc()
        print(f"Pipeline error: {e}")
        yield "", f"⚠️ त्रुटि: {e}", "", None, "❌ *विफल*"


def run_text_pipeline(text_input):
    """Text mode: typed question → answer, citations, audio, status."""
    if not text_input or not text_input.strip():
        yield text_input, "कृपया एक प्रश्न लिखें।", "", None, "🔴 *तैयार*"
        return
    if not HF_TOKEN:
        yield text_input, "⚠️ HF_TOKEN Secret नहीं मिला।", "", None, "❌ *त्रुटि*"
        return

    try:
        yield text_input, "", "", None, "🔎 **जानकारी खोज रहा हूँ...** (RAG)"
        normalized = rag_pipeline.normalize_jargon(text_input)
        eng_query  = rag_pipeline.translate_to_retrieval_query(normalized)
        context_docs = rag_pipeline.retrieve_with_embeddings(eng_query, top_k=3)
        
        citations = "\n".join([f"• {doc['title']}" for doc in context_docs]) if context_docs else ""
        
        yield text_input, "", "", None, "🤖 **उत्तर तैयार कर रहा हूँ...** (LLM)"
        
        context = "\n\n".join(
            f"[{i+1}] {d['title']}\n{d['content'][:500]}"
            for i, d in enumerate(context_docs)
        ) if context_docs else "कोई प्रासंगिक जानकारी नहीं मिली।"

        full_prompt = (
            f"नीचे दी गई जानकारी के आधार पर प्रश्न का उत्तर दें:\n\n"
            f"{context}\n\n"
            f"प्रश्न: {text_input}"
        )
        answer = llm.llm_generate(full_prompt)
        
        yield text_input, answer, citations, None, "🔊 **आवाज़ बना रहा हूँ...** (TTS)"
        clean_text = rag_pipeline.format_response_for_tts(answer)
        audio_out = tts.tts_hindi(clean_text)
        
        yield text_input, answer, citations, audio_out, "✅ **पूरा हुआ**"
        
    except Exception as e:
        import traceback; traceback.print_exc()
        yield text_input, f"⚠️ त्रुटि: {e}", "", None, "❌ *विफल*"


# ─────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────

EXAMPLES = [
    "मुझे होम लोन के लिए कौन से दस्तावेज़ चाहिए?",
    "EMI क्या होती है और कैसे calculate होती है?",
    "मुद्रा लोन कैसे मिलेगा?",
    "CIBIL score क्या होता है?",
    "PM Awas Yojana में subsidy कैसे मिलती है?",
    "जन धन खाता कैसे खोलें?",
    "किसान क्रेडिट कार्ड क्या है?",
    "Atal Pension Yojana में कैसे जुड़ें?",
    "SBI personal loan की interest rate क्या है?",
    "महिला स्वयं सहायता समूह को loan कैसे मिलेगा?",
    "Street vendor ko SVANidhi loan kaise milega?",
    "Bank ke khilaf complaint kahan karein?",
]

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Noto+Sans+Devanagari:wght@300;400;500;700&family=Sora:wght@300;400;600;700&display=swap');

:root {
    --saffron: #FF6B00;
    --saffron-glow: rgba(255,107,0,0.15);
    --green: #0A7A3E;
    --navy: #0F1E3D;
    --navy-mid: #162848;
    --gold: #D4A017;
    --gold-light: #F0C042;
    --text-primary: #F4EFE6;
    --text-secondary: #B8C4D8;
    --border: rgba(212,160,23,0.3);
}
* { box-sizing: border-box; }
body, .gradio-container {
    font-family: 'Sora','Noto Sans Devanagari',sans-serif !important;
    background: var(--navy) !important;
    color: var(--text-primary) !important;
}
.gradio-container { max-width: 1100px !important; margin: 0 auto !important; }

.tricolor { height:4px; background:linear-gradient(90deg,#FF6B00 33%,white 33% 66%,#0A7A3E 66%); }

.app-header {
    background: linear-gradient(135deg,var(--navy) 0%,var(--navy-mid) 60%,#0D2847 100%);
    border-bottom: 2px solid var(--gold);
    padding: 28px 40px 22px; position: relative; overflow: hidden;
}
.app-header::before {
    content:''; position:absolute; top:-60%; right:-5%;
    width:350px; height:350px;
    background:radial-gradient(circle,var(--saffron-glow) 0%,transparent 70%);
    pointer-events:none;
}
.app-header::after { content:'🏦'; position:absolute; right:36px; top:50%; transform:translateY(-50%); font-size:72px; opacity:0.09; }
.header-badge { display:inline-flex; align-items:center; gap:8px; background:rgba(212,160,23,0.1); border:1px solid var(--gold); border-radius:100px; padding:4px 14px; font-size:11px; font-weight:600; letter-spacing:1.5px; text-transform:uppercase; color:var(--gold-light); margin-bottom:10px; }
.header-title { font-size:28px; font-weight:700; color:var(--text-primary); margin-bottom:5px; font-family:'Noto Sans Devanagari','Sora',sans-serif !important; }
.header-subtitle { font-size:13px; color:var(--text-secondary); font-weight:300; }

.notice { background:rgba(212,160,23,0.07); border:1px solid rgba(212,160,23,0.28); border-radius:10px; padding:10px 18px; font-size:12px; color:var(--gold-light); margin:12px 40px; display:flex; align-items:flex-start; gap:10px; line-height:1.6; }

.steps { display:flex; gap:6px; flex-wrap:wrap; padding:12px 40px; background:rgba(10,25,50,0.55); border-bottom:1px solid var(--border); }
.step { display:flex; align-items:center; gap:5px; background:rgba(255,255,255,0.03); border:1px solid var(--border); border-radius:100px; padding:4px 11px; font-size:11px; color:var(--text-secondary); }
.sn { width:16px; height:16px; background:var(--saffron); border-radius:50%; display:flex; align-items:center; justify-content:center; font-size:9px; font-weight:700; color:white; flex-shrink:0; }

label { color:var(--gold-light) !important; font-size:11px !important; font-weight:600 !important; letter-spacing:0.8px !important; text-transform:uppercase !important; }

textarea { background:rgba(10,20,45,0.85) !important; border:1px solid var(--border) !important; border-radius:10px !important; color:var(--text-primary) !important; font-family:'Noto Sans Devanagari','Sora',sans-serif !important; font-size:14px !important; line-height:1.7 !important; padding:13px !important; }
textarea:focus { border-color:var(--saffron) !important; outline:none !important; box-shadow:0 0 0 3px var(--saffron-glow) !important; }

.answer-area textarea { background:linear-gradient(135deg,rgba(10,122,62,0.07) 0%,rgba(10,20,45,0.9) 100%) !important; border:1px solid rgba(10,122,62,0.3) !important; color:#DFF0E6 !important; font-size:15px !important; line-height:1.8 !important; min-height:130px !important; }
.sources-area textarea { background:rgba(212,160,23,0.04) !important; border:1px solid rgba(212,160,23,0.18) !important; font-size:12px !important; color:var(--text-secondary) !important; }

.panel { background:var(--navy-mid) !important; border:1px solid var(--border) !important; border-radius:14px !important; padding:20px !important; margin:8px !important; }

button.primary { background:linear-gradient(135deg,var(--saffron) 0%,#D55000 100%) !important; border:none !important; border-radius:10px !important; color:white !important; font-weight:700 !important; font-size:13px !important; padding:12px 28px !important; cursor:pointer !important; transition:all 0.25s !important; box-shadow:0 4px 20px rgba(255,107,0,0.35) !important; letter-spacing:0.5px !important; text-transform:uppercase !important; }
button.primary:hover { transform:translateY(-2px) !important; box-shadow:0 8px 30px rgba(255,107,0,0.5) !important; }

.tab-nav button { background:transparent !important; border:none !important; border-bottom:2px solid transparent !important; color:var(--text-secondary) !important; font-size:13px !important; font-weight:500 !important; padding:10px 22px !important; border-radius:0 !important; transition:all 0.2s !important; }
.tab-nav button.selected { border-bottom-color:var(--saffron) !important; color:#FF8C33 !important; font-weight:600 !important; }

.footer { padding:16px 40px; text-align:center; font-size:11px; color:rgba(184,196,216,0.4); border-top:1px solid var(--border); background:rgba(8,15,35,0.5); }
::-webkit-scrollbar { width:5px; } ::-webkit-scrollbar-thumb { background:var(--gold); border-radius:3px; }
"""


def build_ui():
    with gr.Blocks(css=CSS, title="बैंकिंग सहायक", theme=gr.themes.Base()) as demo:

        gr.HTML('<div class="tricolor"></div>')
        gr.HTML("""
        <div class="app-header">
            <div class="header-badge">🎙️ Voice · Hindi · Free Tier</div>
            <div class="header-title">बैंकिंग सहायक</div>
            <div class="header-subtitle">AI-powered Hindi Banking & Finance Assistant — Loans · Schemes · EMI · Accounts · Social Security</div>
        </div>
        """)
        gr.HTML("""
        <div class="notice">
            <span>🔑</span>
            <span><b>One-time setup:</b> Space <b>Settings → Secrets</b> → add
            <code style="background:rgba(0,0,0,0.3);padding:1px 6px;border-radius:4px;">HF_TOKEN</code>
            = your free HF token from hf.co/settings/tokens</span>
        </div>
        """)
        gr.HTML("""
        <div class="steps">
            <div class="step"><span class="sn">1</span>हिंदी बोलें</div>
            <div class="step"><span class="sn">2</span>Whisper STT</div>
            <div class="step"><span class="sn">3</span>Jargon Norm</div>
            <div class="step"><span class="sn">4</span>RAG Retrieval</div>
            <div class="step"><span class="sn">5</span>Mistral-7B LLM</div>
            <div class="step"><span class="sn">6</span>MMS Hindi TTS</div>
            <div class="step"><span class="sn">7</span>हिंदी उत्तर</div>
        </div>
        """)

        with gr.Tabs(elem_classes="tab-nav"):

            with gr.Tab("🎙️ Voice Mode"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes="panel"):
                        gr.HTML('<p style="color:#B8C4D8;font-size:13px;margin:0 0 12px;">माइक बटन दबाएं → हिंदी में बोलें → Submit</p>')
                        v_audio_in  = gr.Audio(sources=["microphone"], type="numpy", label="🎤 अपना प्रश्न बोलें")
                        v_btn       = gr.Button("🔍 उत्तर खोजें", variant="primary")
                        v_transcript= gr.Textbox(label="📝 Transcript", interactive=False, lines=2, placeholder="आपकी बात यहाँ दिखेगी...")
                    with gr.Column(scale=1, elem_classes="panel"):
                        v_answer    = gr.Textbox(label="💬 उत्तर", interactive=False, lines=5, placeholder="उत्तर यहाँ आएगा...", elem_classes="answer-area")
                        v_audio_out = gr.Audio(label="🔊 उत्तर सुनें", type="numpy", autoplay=True)
                        v_citations = gr.Textbox(label="📚 स्रोत", interactive=False, lines=3, elem_classes="sources-area")

                v_status = gr.Markdown("🔴 *तैयार*", elem_classes="status-bar")
                v_btn.click(fn=run_voice_pipeline, inputs=[v_audio_in],
                            outputs=[v_transcript, v_answer, v_citations, v_audio_out, v_status])

            with gr.Tab("⌨️ Text Mode"):
                with gr.Row():
                    with gr.Column(scale=1, elem_classes="panel"):
                        gr.HTML('<p style="color:#B8C4D8;font-size:13px;margin:0 0 12px;">हिंदी या Hinglish में टाइप करें</p>')
                        t_input     = gr.Textbox(label="✏️ अपना प्रश्न लिखें", lines=3, placeholder="जैसे: होम लोन के लिए कौन से दस्तावेज़ चाहिए?")
                        t_btn       = gr.Button("🔍 उत्तर खोजें", variant="primary")
                        t_transcript= gr.Textbox(label="📝 आपका प्रश्न", interactive=False, lines=2)
                    with gr.Column(scale=1, elem_classes="panel"):
                        t_answer    = gr.Textbox(label="💬 उत्तर", interactive=False, lines=5, elem_classes="answer-area")
                        t_audio_out = gr.Audio(label="🔊 उत्तर सुनें", type="numpy", autoplay=True)
                        t_citations = gr.Textbox(label="📚 स्रोत", interactive=False, lines=3, elem_classes="sources-area")

                t_status = gr.Markdown("🔴 *तैयार*", elem_classes="status-bar")
                t_btn.click(fn=run_text_pipeline, inputs=[t_input],
                            outputs=[t_transcript, t_answer, t_citations, t_audio_out, t_status])
                t_input.submit(fn=run_text_pipeline, inputs=[t_input],
                               outputs=[t_transcript, t_answer, t_citations, t_audio_out, t_status])

        gr.HTML('<div style="padding:6px 10px 0;"><p style="color:var(--gold);font-size:11px;font-weight:600;letter-spacing:1px;text-transform:uppercase;margin:8px 0 4px 32px;">📌 उदाहरण प्रश्न</p></div>')
        gr.Examples(examples=[[q] for q in EXAMPLES], inputs=[t_input], label="", cache_examples=False)

        with gr.Accordion("ℹ️ About / जानकारी", open=False):
            gr.HTML("""
            <div style="padding:18px;color:#B8C4D8;font-size:13px;line-height:1.8;">
              <div style="display:grid;grid-template-columns:1fr 1fr;gap:20px;margin-bottom:14px;">
                <div>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:8px;">🤖 Models</div>
                  <div>• <b>STT:</b> openai/whisper-large-v3 (hf-inference)</div>
                  <div>• <b>LLM:</b> mistralai/Mistral-7B-Instruct-v0.3 (featherless)</div>
                  <div>• <b>TTS:</b> facebook/mms-tts-hin (hf-inference)</div>
                  <br/>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:6px;">⏱️ Response Time</div>
                  <div>• STT: ~5–10s &nbsp;• LLM: ~5–15s &nbsp;• TTS: ~3–8s</div>
                  <div><b>Total: ~15–35s per query</b></div>
                </div>
                <div>
                  <div style="color:#D4A017;font-weight:600;margin-bottom:8px;">📖 Knowledge Base</div>
                  <div>• 12 Major Indian Banks</div>
                  <div>• 11 Govt Credit Schemes (Mudra, PMAY, KCC...)</div>
                  <div>• 7 Social Security Schemes (APY, PMJJBY, EPFO...)</div>
                  <div>• Financial Terms (EMI, CIBIL, KYC, MCLR...)</div>
                  <div>• MFI, NBFC, RRBs, Co-op Banks</div>
                  <div>• Consumer Rights & RBI Ombudsman</div>
                </div>
              </div>
              <div style="background:rgba(255,107,0,0.07);border:1px solid rgba(255,107,0,0.2);border-radius:8px;padding:10px 14px;font-size:12px;color:#FF8C33;">
                ⚠️ यह assistant केवल सामान्य जानकारी देता है। वित्तीय निर्णयों के लिए अपने बैंक या certified advisor से सलाह लें।
              </div>
            </div>
            """)

        gr.HTML('<div class="footer">बैंकिंग सहायक · Whisper + Mistral-7B + MMS-TTS · CPU Basic Safe · HF Inference Providers</div>')

    return demo


if __name__ == "__main__":
    demo = build_ui()
    demo.queue().launch(show_api=False)