src/streamlit_app.py

import streamlit as st
import os, requests, hmac, re, json, base64
from concurrent.futures import ThreadPoolExecutor, as_completed

st.set_page_config(page_title="Yai", page_icon="✦", layout="wide", initial_sidebar_state="collapsed")

st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Google+Sans:wght@300;400;500&display=swap');
html,body,[data-testid="stApp"]{background:#0f0f11!important;color:#e3e3e6!important;font-family:'Google Sans',sans-serif!important;}
#MainMenu,footer,header,[data-testid="stToolbar"],[data-testid="stDecoration"],[data-testid="stStatusWidget"]{display:none!important;}
.main .block-container{max-width:760px!important;margin:0 auto!important;padding:0 16px 140px!important;}
[data-testid="stSidebar"]{background:#1a1a1f!important;border-right:1px solid #2a2a32!important;}
[data-testid="stSidebar"] *{color:#b4b4be!important;}

/* ── Logo ── */
.yai-wrap{display:flex;flex-direction:column;align-items:center;padding:48px 0 28px;gap:10px;}
.yai-icon{width:68px;height:68px;background:linear-gradient(135deg,#4f8ef7,#a259f7,#f759ab);border-radius:22px;display:flex;align-items:center;justify-content:center;font-size:34px;font-weight:700;color:#fff;box-shadow:0 8px 32px rgba(79,142,247,.35);}
.yai-name{font-size:38px;font-weight:300;background:linear-gradient(135deg,#4f8ef7,#a259f7,#f759ab);-webkit-background-clip:text;-webkit-text-fill-color:transparent;}
.yai-sub{font-size:14px;color:#555566;text-align:center;}

/* ── Suggestion cards ── */
.sugg-grid{display:grid;grid-template-columns:1fr 1fr;gap:10px;margin:20px 0 28px;}
.sugg-card{background:#16161e;border:1px solid #24242e;border-radius:14px;padding:14px;cursor:pointer;transition:all .2s;}
.sugg-card:hover{border-color:#4f8ef7;background:rgba(79,142,247,.05);}
.sugg-icon{font-size:18px;margin-bottom:6px;}
.sugg-title{font-size:13px;font-weight:500;color:#d0d0da;}
.sugg-sub{font-size:11px;color:#4a4a58;margin-top:3px;}

/* ── Mode pills row ── */
.mode-row{display:flex;gap:6px;flex-wrap:wrap;margin-bottom:14px;}
.mpill{display:inline-flex;align-items:center;gap:5px;padding:6px 14px;border-radius:50px;border:1px solid #2a2a35;background:#16161e;color:#8a8a9a;font-size:12px;cursor:pointer;transition:all .2s;}
.mpill.active,.mpill:hover{border-color:#4f8ef7;color:#4f8ef7;background:rgba(79,142,247,.08);}

/* ── Messages – Gemini style: no bubbles ── */
[data-testid="stChatMessage"]{background:transparent!important;border:none!important;padding:18px 0 4px!important;}
[data-testid="stChatMessageContent"]{background:transparent!important;border:none!important;padding:0!important;font-size:15px;line-height:1.7;}
[data-testid="stChatMessageAvatar"]{width:30px!important;height:30px!important;min-width:30px!important;border-radius:50%!important;background:linear-gradient(135deg,#4f8ef7,#a259f7)!important;}

/* User msg right-align */
[data-testid="stChatMessage"]:has([alt="user"]){flex-direction:row-reverse!important;}
[data-testid="stChatMessage"]:has([alt="user"]) [data-testid="stChatMessageContent"]{text-align:right;color:#c8d8ff;}

/* ── Action icons below msg ── */
.msg-actions{display:flex;gap:14px;padding:6px 0 2px;opacity:.5;transition:opacity .2s;}
.msg-actions:hover{opacity:1;}
.msg-actions span{font-size:15px;cursor:pointer;}

/* ── Code blocks ── */
pre{background:#12121a!important;border:1px solid #2a2a38!important;border-radius:10px!important;padding:14px!important;font-size:13px!important;overflow-x:auto!important;}
code{background:#1e1e28!important;padding:2px 5px!important;border-radius:4px!important;font-size:12px!important;color:#a8d8a8!important;}

/* ── Reasoning expander ── */
[data-testid="stExpander"]{background:#13131a!important;border:1px solid #2a2a35!important;border-radius:10px!important;margin-bottom:10px!important;}
[data-testid="stExpander"] summary{color:#555568!important;font-size:12px!important;}

/* ── Input bar ── */
[data-testid="stBottom"]{background:#0f0f11!important;border-top:1px solid #1a1a24!important;padding:10px 0!important;}
.stChatInputContainer{background:#1a1a1f!important;border:1px solid #2e2e3a!important;border-radius:26px!important;padding:4px 8px 4px 16px!important;max-width:760px!important;margin:0 auto!important;box-shadow:0 4px 20px rgba(0,0,0,.4)!important;}
.stChatInputContainer:focus-within{border-color:#4f8ef7!important;}
.stChatInputContainer textarea{background:transparent!important;border:none!important;color:#e3e3e6!important;font-family:'Google Sans',sans-serif!important;font-size:15px!important;}
.stChatInputContainer textarea::placeholder{color:#3a3a4a!important;}
.stChatInputContainer button{background:linear-gradient(135deg,#4f8ef7,#a259f7)!important;border:none!important;border-radius:50%!important;width:34px!important;height:34px!important;}

/* ── Pipeline steps ── */
.pstep{display:flex;align-items:center;gap:8px;padding:7px 12px;background:#13131c;border:1px solid #22222c;border-radius:8px;margin:3px 0;font-size:12px;color:#666676;animation:fi .3s ease;}
.pstep.ok{border-color:#1a3a1a;color:#3a8a3a;background:#0d160d;}
.dot{width:5px;height:5px;border-radius:50%;background:#4f8ef7;animation:pu 1s infinite;}
.pstep.ok .dot{background:#3fb950;animation:none;}
@keyframes pu{0%,100%{opacity:1}50%{opacity:.2}}
@keyframes fi{from{opacity:0;transform:translateY(3px)}to{opacity:1;transform:translateY(0)}}

/* ── Model tag ── */
.mtag{display:inline-flex;align-items:center;gap:4px;padding:2px 9px;background:#1a1a28;border:1px solid #2a2a3a;border-radius:20px;font-size:10px;color:#5a6af0;margin-top:5px;}

/* ── Buttons ── */
.stButton button{background:#1a1a22!important;color:#9a9aaa!important;border:1px solid #2a2a35!important;border-radius:8px!important;font-size:12px!important;transition:all .2s!important;}
.stButton button:hover{border-color:#4f8ef7!important;color:#4f8ef7!important;}
hr{border-color:#1e1e28!important;}
::-webkit-scrollbar{width:4px;}::-webkit-scrollbar-thumb{background:#2a2a38;border-radius:2px;}
</style>
""", unsafe_allow_html=True)

# ── Constants ──────────────────────────────────────────
HF_TOKEN = os.getenv("HF_TOKEN","")
JINA_KEY = os.getenv("JINA_KEY","")
PASSWORD = "132733yy"
HF_API   = "https://router.huggingface.co/v1/chat/completions"

MODELS = {
    "ceo":        "deepseek-ai/DeepSeek-V3-0324:novita",
    "researcher": "moonshotai/Kimi-K2-Instruct:novita",
    "coder":      "Qwen/Qwen3-Coder-480B-A35B-Instruct:novita",
    "tutor":      "meta-llama/Llama-3.1-8B-Instruct:cerebras",
    "reasoner":   "deepseek-ai/DeepSeek-R1:novita",
}
INFO = {
    "ceo":        {"icon":"✦",  "label":"Orchestrator","model":"DeepSeek V3"},
    "researcher": {"icon":"🌐", "label":"Researcher",  "model":"Kimi K2"},
    "coder":      {"icon":"⌨️", "label":"Coder",       "model":"Qwen3-Coder"},
    "tutor":      {"icon":"📖", "label":"Tutor",       "model":"Llama 3.1"},
    "reasoner":   {"icon":"🧠", "label":"Reasoner",    "model":"DeepSeek R1"},
}
PROMPTS = {
    "ceo":"""You are an intelligent orchestrator. Read the user message and decide which agents to use.
Reply ONLY with this JSON (no extra text):
{"plan":"one sentence","agents":[{"name":"agent","task":"specific task"}]}
Rules: simple question→tutor only | code→coder | URL→researcher | math/logic→reasoner | max 3 agents.
Always respond in the same language as the user.""",
    "researcher":"""You are a precise research analyst. Respond in the user's language.
Analyze URLs and content deeply. Cite sources. Say "I'm not sure" when needed.
Format: key points → analysis → your assessment. Never fabricate.""",
    "coder":"""You are an expert software engineer. Respond in the user's language.
Write clean, documented, working code. Always include error handling and usage examples.
Warn about bugs immediately. Ask for clarification if ambiguous.""",
    "tutor":"""You are an exceptional educator. Respond in the user's language (Arabic→Arabic, English→English).
Start with a real-world analogy, then explain from simple to complex.
Never fabricate. End with a check question.""",
    "reasoner":"""You are a deep reasoning AI. Respond in the user's language.
Solve complex problems step by step. Show your reasoning clearly.
Compare solutions when multiple exist. Accuracy over speed.""",
}
MODES = [("💬","Chat"),("🔍","Search"),("⌨️","Code"),("🎨","Image"),("🎙️","Voice"),("🎬","Video")]

# ── Login ──────────────────────────────────────────────
def check_password():
    if st.session_state.get("auth"): return True
    st.markdown('<div class="yai-wrap"><div class="yai-icon">Y</div><div class="yai-name">Yai</div><div class="yai-sub">Your personal AI</div></div>',unsafe_allow_html=True)
    c1,c2,c3=st.columns([1,1.2,1])
    with c2:
        pwd=st.text_input("",placeholder="Password...",type="password",key="pw",label_visibility="collapsed")
        if st.button("Continue →",use_container_width=True):
            if hmac.compare_digest(pwd,PASSWORD):
                st.session_state["auth"]=True; st.rerun()
            else: st.error("Incorrect password")
    return False
if not check_password(): st.stop()

# ── Helpers ────────────────────────────────────────────
def fetch_url(url):
    try:
        h={"Accept":"text/plain"}
        if JINA_KEY: h["Authorization"]=f"Bearer {JINA_KEY}"
        r=requests.get(f"https://r.jina.ai/{url}",headers=h,timeout=15)
        return r.text[:4000] if r.status_code==200 else f"[HTTP {r.status_code}]"
    except Exception as e: return f"[Error:{e}]"

def extract_urls(t): return re.findall(r'https?://[^\s]+',t)

def call_llm(messages,model_key,max_tokens=1200):
    if not HF_TOKEN: return "⚠️ HF_TOKEN missing in Space secrets."
    try:
        r=requests.post(HF_API,
            headers={"Authorization":f"Bearer {HF_TOKEN}","Content-Type":"application/json"},
            json={"model":MODELS[model_key],"messages":messages,"max_tokens":max_tokens,"temperature":0.7,"stream":False},
            timeout=90)
        r.raise_for_status()
        return r.json()["choices"][0]["message"]["content"]
    except requests.exceptions.HTTPError:
        c=r.status_code
        if c==503: return "⏳ Model loading, retry in a moment."
        if c==429: return "🚦 Rate limit, please wait."
        return f"❌ HTTP {c}: {r.text[:200]}"
    except Exception as e: return f"❌ Error: {e}"

def run_agent(name,task,ctx):
    msgs=[{"role":"system","content":PROMPTS[name]},{"role":"user","content":f"Task: {task}\n\nContext:\n{ctx}"}]
    return name, call_llm(msgs,name)

def generate_image(prompt):
    try:
        r=requests.post("https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell",
            headers={"Authorization":f"Bearer {HF_TOKEN}"},
            json={"inputs":prompt},timeout=60)
        if r.status_code==200:
            return "data:image/png;base64,"+base64.b64encode(r.content).decode()
    except: pass
    return None

def generate_voice(text):
    try:
        r=requests.post("https://api-inference.huggingface.co/models/microsoft/speecht5_tts",
            headers={"Authorization":f"Bearer {HF_TOKEN}"},
            json={"inputs":text[:400]},timeout=60)
        if r.status_code==200: return r.content
    except: pass
    return None

# ── Pipeline ───────────────────────────────────────────
def pipeline(user_msg, history, sbox, mode):
    enriched=user_msg
    urls=extract_urls(user_msg)
    if urls:
        with sbox: st.markdown('<div class="pstep"><span class="dot"></span>Reading URLs with Jina...</div>',unsafe_allow_html=True)
        for u in urls[:2]: enriched+=f"\n\n[URL:{u}]\n{fetch_url(u)}"

    if mode=="image":
        with sbox: st.markdown('<div class="pstep"><span class="dot"></span>Generating image with FLUX.1...</div>',unsafe_allow_html=True)
        img=generate_image(user_msg); sbox.empty()
        return ("IMG",img) if img else ("TXT","⚠️ Image generation failed — model may be loading. Try again in 30s."), [], {}

    if mode=="voice":
        with sbox: st.markdown('<div class="pstep"><span class="dot"></span>Generating voice...</div>',unsafe_allow_html=True)
        audio=generate_voice(user_msg); sbox.empty()
        return (("AUDIO",audio) if audio else ("TXT","⚠️ Voice generation failed.")), [], {}

    if mode=="video":
        sbox.empty()
        return ("TXT","🎬 For video generation, use [Kling AI](https://klingai.com) — free & high quality. Open in new tab, describe your scene, download the result."), [], {}

    # CEO plan
    with sbox: st.markdown('<div class="pstep"><span class="dot"></span>Orchestrator analyzing...</div>',unsafe_allow_html=True)
    hist_txt="\n".join([f"{m['role']}: {m['content'][:150]}" for m in history[-6:]])
    ceo_msgs=[{"role":"system","content":PROMPTS["ceo"]},{"role":"user","content":f"History:\n{hist_txt}\n\nMessage:\n{enriched}"}]
    raw=call_llm(ceo_msgs,"ceo",512)

    try:
        m=re.search(r'\{.*\}',raw,re.DOTALL)
        d=json.loads(m.group()) if m else {}
        agents=d.get("agents",[{"name":"tutor","task":enriched}])
    except: agents=[{"name":"tutor","task":enriched}]

    valid=set(MODELS)-{"ceo"}
    agents=[a for a in agents if a.get("name") in valid] or [{"name":"tutor","task":enriched}]

    with sbox:
        for a in agents:
            i=INFO[a["name"]]
            st.markdown(f'<div class="pstep"><span class="dot"></span>{i["icon"]} {i["label"]} — {a["task"][:50]}...</div>',unsafe_allow_html=True)

    results={}
    with ThreadPoolExecutor(max_workers=len(agents)) as ex:
        futs={ex.submit(run_agent,a["name"],a["task"],enriched):a["name"] for a in agents}
        for f in as_completed(futs):
            n,r=f.result(); results[n]=r
            with sbox: st.markdown(f'<div class="pstep ok"><span class="dot"></span>{INFO[n]["icon"]} {INFO[n]["label"]} done</div>',unsafe_allow_html=True)

    if len(results)==1:
        return ("TXT",list(results.values())[0]), agents, results

    with sbox: st.markdown('<div class="pstep"><span class="dot"></span>Synthesizing final answer...</div>',unsafe_allow_html=True)
    synth=[{"role":"system","content":"You are an orchestrator synthesizing team results. Respond in the user's language. Merge insights into one cohesive, well-structured response."},
           {"role":"user","content":f"Question:\n{user_msg}\n\nTeam:\n"+ "\n\n".join([f"=={INFO[n]['label']}==\n{r}" for n,r in results.items()])+"\n\nWrite unified final response."}]
    final=call_llm(synth,"ceo",1200)
    return ("TXT",final), agents, results

# ── Session ────────────────────────────────────────────
for k,v in [("msgs",[]),("mode","chat"),("upload_ctx","")]:
    if k not in st.session_state: st.session_state[k]=v

# ── Sidebar ────────────────────────────────────────────
with st.sidebar:
    st.markdown("### ✦ Yai"); st.divider()
    st.markdown("**Active models:**")
    for k,i in INFO.items(): st.markdown(f"{i['icon']} `{i['model']}`")
    st.divider()
    st.markdown(f"HF Token: {'🟢' if HF_TOKEN else '🔴'}")
    st.markdown(f"Jina: {'🟢' if JINA_KEY else '🟡'}")
    st.divider()
    show_think=st.toggle("Show reasoning",True)
    show_agents=st.toggle("Show agent details",False)
    st.divider()
    if st.button("🗑️ Clear chat",use_container_width=True):
        st.session_state.msgs=[]; st.rerun()

# ── Main ───────────────────────────────────────────────
if not st.session_state.msgs:
    st.markdown("""
    <div class="yai-wrap">
        <div class="yai-icon">Y</div>
        <div class="yai-name">Yai</div>
        <div class="yai-sub">Ask anything — 5 AI models working together</div>
    </div>
    <div class="sugg-grid">
        <div class="sugg-card"><div class="sugg-icon">💡</div><div class="sugg-title">Explain a concept</div><div class="sugg-sub">Any topic, any level</div></div>
        <div class="sugg-card"><div class="sugg-icon">⌨️</div><div class="sugg-title">Write code</div><div class="sugg-sub">Any language or framework</div></div>
        <div class="sugg-card"><div class="sugg-icon">🔍</div><div class="sugg-title">Research a topic</div><div class="sugg-sub">Paste a URL or ask</div></div>
        <div class="sugg-card"><div class="sugg-icon">🧠</div><div class="sugg-title">Solve a problem</div><div class="sugg-sub">Logic, math, analysis</div></div>
    </div>""", unsafe_allow_html=True)

# Mode bar
cols=st.columns(len(MODES))
for i,(icon,label) in enumerate(MODES):
    with cols[i]:
        active=st.session_state.mode==label.lower()
        if st.button(f"{icon} {label}",key=f"m{i}",type="primary" if active else "secondary",use_container_width=True):
            st.session_state.mode=label.lower(); st.rerun()

# File upload
c1,c2=st.columns([8,1])
with c2:
    if st.button("📎",help="Attach file"):
        st.session_state["show_up"]=not st.session_state.get("show_up",False); st.rerun()
if st.session_state.get("show_up"):
    up=st.file_uploader("",type=["txt","py","js","md","csv","png","jpg","pdf"],label_visibility="collapsed")
    if up:
        if up.type.startswith("image/"):
            b64=base64.b64encode(up.read()).decode()
            st.session_state.upload_ctx=f"[Image: {up.name}]"
            st.image(f"data:{up.type};base64,{b64}",width=180)
        else:
            try:
                content=up.read().decode("utf-8",errors="ignore")
                st.session_state.upload_ctx=f"[File: {up.name}]\n{content[:3000]}"
                st.success(f"✓ {up.name}")
            except: st.error("Could not read file")
st.markdown("---")

# Chat history
for msg in st.session_state.msgs:
    if msg["role"]=="user":
        with st.chat_message("user"):
            st.markdown(msg["content"])
    else:
        with st.chat_message("assistant"):
            rt=msg.get("rt","TXT")
            if rt=="IMG": st.image(msg["content"],use_container_width=True)
            elif rt=="AUDIO": st.audio(msg["content"],format="audio/wav")
            else:
                content=msg["content"]
                think=re.search(r'<think>(.*?)</think>',content,re.DOTALL)
                if think and show_think:
                    with st.expander("🧠 View reasoning"):
                        st.markdown(f"```\n{think.group(1).strip()}\n```")
                clean=re.sub(r'<think>.*?</think>','',content,flags=re.DOTALL).strip()
                st.markdown(clean)
            st.markdown('<div class="msg-actions"><span title="Copy">⎘</span><span title="Retry">↺</span><span title="Good">👍</span><span title="Bad">👎</span></div>',unsafe_allow_html=True)
            if show_agents and msg.get("details") and len(msg["details"])>1:
                with st.expander("View agent responses"):
                    for n,r in msg["details"].items():
                        i=INFO.get(n,{}); st.markdown(f"**{i.get('icon','')} {i.get('label',n)}** · `{i.get('model','')}`")
                        st.markdown(r); st.divider()
            if msg.get("agents"):
                tags=" · ".join([INFO[a["name"]]["model"] for a in msg["agents"] if a["name"] in INFO])
                st.markdown(f'<div class="mtag">✦ {tags}</div>',unsafe_allow_html=True)

# Input
pmap={"chat":"Ask Yai anything...","search":"Paste a URL or ask...","code":"Describe the code you need...","image":"Describe the image...","voice":"Text to convert to voice...","video":"Describe the video scene..."}
if user_input:=st.chat_input(pmap.get(st.session_state.mode,"Ask Yai...")):
    full=user_input
    if st.session_state.upload_ctx:
        full+="\n\n"+st.session_state.upload_ctx
        st.session_state.upload_ctx=""
        st.session_state["show_up"]=False

    st.session_state.msgs.append({"role":"user","content":user_input})
    with st.chat_message("user"): st.markdown(user_input)

    with st.chat_message("assistant"):
        sbox=st.container()
        history=[m for m in st.session_state.msgs if m["role"]!="system"][:-1]
        result,agents,details=pipeline(full,history,sbox,st.session_state.mode)
        sbox.empty()
        rt,rc=result

        if rt=="IMG": st.image(rc,use_container_width=True)