import gradio as gr
import pandas as pd
import os
import time
from groq import Groq
from difflib import SequenceMatcher

# API Key Configuration
GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
client = Groq(api_key=GROQ_API_KEY) if GROQ_API_KEY else None

# --- ၁။ Dataset Loading (Safe Logic) ---
chat_db = []
try:
    csv_url = "https://huggingface.co/datasets/amkyawdev/AmkyawDev-Dataset/raw/main/train.csv"
    df = pd.read_csv(csv_url)
    if not df.empty:
        for _, row in df.iterrows():
            if len(row) >= 2:
                u_text = str(row.iloc[0]).strip()
                a_text = str(row.iloc[1]).strip()
                if u_text.lower() != "nan" and a_text.lower() != "nan":
                    chat_db.append({"u": u_text, "a": a_text})
    print(f"✅ Dataset Loaded: {len(chat_db)} pairs.")
except Exception as e:
    print(f"⚠️ Dataset Warning: {e}")

# --- ၂။ AI Logic (Llama-3.1-8b-instant) ---
def generate_response(user_input, history):
    user_input = user_input.strip()

    # Step 1: CSV Matching (Fuzzy Search)
    for pair in chat_db:
        if SequenceMatcher(None, user_input.lower(), pair["u"].lower()).ratio() > 0.85:
            res = pair["a"]
            for i in range(len(res)):
                time.sleep(0.005)
                yield res[: i + 1]
            return

    # Step 2: Groq API Call
    if client:
        try:
            # Gradio 4 ရဲ့ [[u, a]] format ကို Groq ရဲ့ role format ပြောင်းခြင်း
            messages = [{"role": "system", "content": "You are Amkyaw AI, a friendly Myanmar assistant. Use Unicode."}]
            for h in history:
                if h[0]: messages.append({"role": "user", "content": h[0]})
                if h[1]: messages.append({"role": "assistant", "content": h[1]})
            messages.append({"role": "user", "content": user_input})

            stream = client.chat.completions.create(
                messages=messages,
                model="llama-3.1-8b-instant",
                temperature=0.7,
                stream=True,
            )
            
            full_res = ""
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    full_res += chunk.choices[0].delta.content
                    yield full_res
            return
        except Exception as e:
            print(f"Groq API Error: {e}")

    yield "တောင်းပန်ပါတယ်၊ အခုလောလောဆယ် မဖြေနိုင်သေးပါဘူးခင်ဗျာ။"

# --- ၃။ UI Design (Gradio 4 Style) ---
custom_css = """
footer { visibility: hidden !important; }
.gradio-container { background: #0d1117 !important; color: white !important; font-family: sans-serif; }
.header-card {
    background: rgba(255, 255, 255, 0.05);
    border: 1px solid rgba(255, 255, 255, 0.1);
    border-radius: 12px; padding: 15px; margin-bottom: 20px;
}
#input-row {
    background: #161b22 !important; border-radius: 25px !important;
    border: 1px solid #30363d !important; padding: 5px 15px !important;
}
#send-btn {
    background: #ffffff !important; color: #000000 !important;
    border-radius: 50% !important; width: 40px !important; height: 40px !important; min-width: 40px !important;
}
"""

with gr.Blocks(title="Amkyaw AI v4", css=custom_css) as demo:
    gr.HTML("""
        <div class="header-card">
            <div style="display:flex; justify-content: space-between; align-items: center;">
                <b style="color: #58a6ff; font-size: 20px;">Amkyaw AI</b>
                <span style="color:#8b949e; font-size:11px;">Groq | Gradio 4 Stable</span>
            </div>
        </div>
    """)
    
    chatbot = gr.Chatbot(show_label=False, height=600)
    
    with gr.Row(elem_id="input-row"):
        msg = gr.Textbox(placeholder="Ask me anything...", show_label=False, scale=15, container=False)
        submit = gr.Button("↑", elem_id="send-btn", scale=1)

    # UI Flow logic
    def user_flow(user_msg, history):
        if not user_msg: return "", history
        return "", history + [[user_msg, None]]

    def bot_flow(history):
        user_input = history[-1][0]
        # Streaming response back to chatbot
        for chunk in generate_response(user_input, history[:-1]):
            history[-1][1] = chunk
            yield history

    # Event binding
    msg.submit(user_flow, [msg, chatbot], [msg, chatbot], queue=False).then(bot_flow, chatbot, chatbot)
    submit.click(user_flow, [msg, chatbot], [msg, chatbot], queue=False).then(bot_flow, chatbot, chatbot)

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