# app.py (robust - paste entire file)
import streamlit as st
import torch
import transformers

# robust pipeline import
try:
    from transformers import pipeline
except Exception:
    from transformers.pipelines import pipeline  # fallback

st.set_page_config(page_title="GUVI AI Assistant", page_icon="🤖", layout="wide")

@st.cache_resource
def load_models():
    try:
        translator = pipeline("translation", model="facebook/mbart-large-50-many-to-many-mmt",
                              device=0 if torch.cuda.is_available() else -1)
        generator = pipeline("text-generation", model="aswinprasath31/guvi-gpt2-finetuned",
                             device=0 if torch.cuda.is_available() else -1, return_full_text=False)
        detector = pipeline("text-classification", model="papluca/xlm-roberta-base-language-detection",
                            device=0 if torch.cuda.is_available() else -1)
    except Exception as e:
        st.error(f"Model load failed: {e}")
        raise
    return translator, generator, detector

with st.spinner("Initializing AI models, please wait..."):
    translator, generator, detector = load_models()

MODEL_LANG_MAP = {
    "ar": "ar_AR", "cs": "cs_CZ", "de": "de_DE", "en": "en_XX", "es": "es_XX", "fr": "fr_XX",
    "gu": "gu_IN", "hi": "hi_IN", "it": "it_IT", "ja": "ja_XX", "ko": "ko_KR", "ml": "ml_IN",
    "mr": "mr_IN", "pt": "pt_XX", "ta": "ta_IN", "te": "te_IN", "zh": "zh_CN"
}

st.sidebar.title("About the Chatbot")
st.sidebar.info("Advanced multilingual AI assistant for GUVI (fine-tuned GPT-2).")
st.title("GUVI AI Learning Assistant 🤖")

if "messages" not in st.session_state:
    st.session_state.messages = [{"role":"assistant","content":"Hello! How can I help you with GUVI's courses today?"}]
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

SYSTEM_PROMPT = (
    "You are GUVI Assistant — helpful, factual, concise. "
    "If unsure, reply: 'I don't know — please check GUVI's site or ask support.' "
    "Do NOT invent statistics or promotional claims."
)
FEW_SHOT = ("Q: What is Data Science?\nA: Data Science is extracting insights from data "
           "using statistics, ML, and visualization.\n---\n")

def build_prompt(q): return f"{SYSTEM_PROMPT}\n\n{FEW_SHOT}\nUser: {q}\nAssistant:"

def clean_output(text):
    if not text: return text
    if "Assistant:" in text: text = text.split("Assistant:")[-1]
    lines = [l.strip() for l in text.splitlines() if l.strip()]
    if not lines: return text.strip()
    cleaned = [lines[0]]
    for ln in lines[1:]:
        if ln != cleaned[-1]:
            cleaned.append(ln)
    for p in ["we have taken","GEEKOINS","for FREE","over 2,000"]:
        cleaned = [ln.replace(p, "") for ln in cleaned]
    return "\n".join([ln for ln in cleaned if ln]).strip()

def get_bot_response(user_input):
    if len(user_input.split()) < 3:
        return "Please ask a more specific question."
    try:
        det = detector(user_input)
        label = det[0].get("label","").lower()
        lang2 = label[:2] if label and label.isalpha() else "en"
    except Exception:
        lang2 = "en"
    src = MODEL_LANG_MAP.get(lang2, "en_XX")
    if src == "en_XX":
        en_q = user_input
    else:
        try:
            en_q = translator(user_input, src_lang=src, tgt_lang="en_XX")[0].get("translation_text", user_input)
        except Exception:
            en_q = user_input
    prompt = build_prompt(en_q)
    try:
        out = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.22,
                        top_p=0.92, top_k=50, repetition_penalty=1.2, num_return_sequences=1)
        if isinstance(out, list) and isinstance(out[0], dict):
            en_resp = out[0].get("generated_text") or out[0].get("text","")
        else:
            en_resp = str(out[0]) if isinstance(out, list) else str(out)
    except Exception as e:
        st.error(f"Generation error: {e}")
        return "Sorry, generation failed."
    en_resp = clean_output(en_resp)
    if src == "en_XX":
        return en_resp
    try:
        return translator(en_resp, src_lang="en_XX", tgt_lang=src)[0].get("translation_text", en_resp)
    except Exception:
        return en_resp

if prompt := st.chat_input("Ask me about Data Science, Python, or other courses..."):
    st.session_state.messages.append({"role":"user","content":prompt})
    with st.chat_message("user"): st.markdown(prompt)
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            resp = get_bot_response(prompt)
        st.markdown(resp)
    st.session_state.messages.append({"role":"assistant","content":resp})

st.markdown("---")
st.markdown('<div style="text-align:center;color:grey;">Developed by <b>Aswinprasath V</b></div>', unsafe_allow_html=True)