# -----------------------------------------------------------
# FINAL WORKING app.py FOR HUGGINGFACE SPACE (MIC + AVATAR)
# -----------------------------------------------------------

import os
import uuid
import tempfile
import json
from pathlib import Path
from datetime import datetime

import gradio as gr

# ---------------- STT SETUP ----------------
WHISPER_AVAILABLE = False
try:
    from faster_whisper import WhisperModel
    whisper_model = WhisperModel("small", device="cpu", compute_type="int8")
    WHISPER_AVAILABLE = True
except Exception:
    try:
        import whisper
        whisper_model = whisper.load_model("small")
    except Exception:
        whisper_model = None

# ---------------- LLM SETUP ----------------
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

LLM_MODEL = "google/flan-t5-small"
try:
    tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
    model = AutoModelForSeq2SeqLM.from_pretrained(LLM_MODEL)
    llm_pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
except:
    llm_pipe = None

# ---------------- TTS ----------------
from gtts import gTTS
from pydub import AudioSegment

# ---------------- MEMORY ----------------
MEMORY_FILE = Path("memory.json")
if MEMORY_FILE.exists():
    try:
        with open(MEMORY_FILE, "r", encoding="utf-8") as f:
            MEMORY = json.load(f)
    except:
        MEMORY = {"history": []}
else:
    MEMORY = {"history": []}

def save_memory():
    with open(MEMORY_FILE, "w", encoding="utf-8") as f:
        json.dump(MEMORY, f, ensure_ascii=False, indent=2)

# -----------------------------------------------------------
# ---------------------- HELPERS -----------------------------
# -----------------------------------------------------------

def transcribe_audio(path):
    """Use whisper or fast-whisper"""
    if whisper_model is None:
        return ""
    try:
        if WHISPER_AVAILABLE:
            segments, _ = whisper_model.transcribe(path)
            return " ".join([s.text for s in segments])
        else:
            res = whisper_model.transcribe(path)
            return res.get("text", "")
    except:
        return ""

def ask_llm(text):
    """LLM + memory context"""
    if llm_pipe is None:
        return "I am having trouble thinking right now."

    recent = MEMORY.get("history", [])[-6:]
    ctx = ""
    for m in recent:
        if "user" in m:
            ctx += f"User: {m['user']}\n"
        if "assistant" in m:
            ctx += f"Assistant: {m['assistant']}\n"

    prompt = f"""
You are a smart, warm, concise AI companion.
Use the user's language. Be clear. Give one follow-up question.

{ctx}
User: {text}
Assistant:
    """

    out = llm_pipe(prompt, max_length=200, do_sample=False)
    ans = out[0]["generated_text"]
    if "Assistant:" in ans:
        ans = ans.split("Assistant:")[-1].strip()
    return ans.strip()

def tts(text, out_path, lang="en"):
    try:
        if not text.strip():
            silent = AudioSegment.silent(duration=300)
            silent.export(out_path, format="mp3")
            return out_path
        gTTS(text=text, lang=lang).save(out_path)
        return out_path
    except:
        silent = AudioSegment.silent(duration=300)
        silent.export(out_path, format="mp3")
        return out_path

# -----------------------------------------------------------
# -------------- INTERACTION PIPELINE -----------------------
# -----------------------------------------------------------

def process(avatar_upload, mic_audio, typed_text, tts_lang):
    uid = uuid.uuid4().hex[:8]
    tmpdir = Path(tempfile.gettempdir()) / f"a_{uid}"
    tmpdir.mkdir(exist_ok=True)

    # Avatar selection
    repo_avatar = Path("avatar.png")
    if repo_avatar.exists():
        avatar_path = "avatar.png"
    elif avatar_upload:
        ap = tmpdir / "avatar.png"
        with open(ap, "wb") as f:
            f.write(avatar_upload.read())
        avatar_path = str(ap)
    else:
        avatar_path = None

    if avatar_path is None:
        return "<b>No avatar.png found</b>", "No avatar", None

    # Determine user text
    if typed_text and typed_text.strip():
        user_text = typed_text.strip()
    else:
        if mic_audio is None:
            return None, "Speak or type something!", None
        wav = tmpdir / "mic.wav"
        with open(wav, "wb") as f:
            f.write(mic_audio.read())
        user_text = transcribe_audio(str(wav))

    if not user_text:
        return None, "Didn't catch that. Try again.", None

    # Save to memory
    MEMORY["history"].append({"user": user_text})
    MEMORY["history"] = MEMORY["history"][-200:]
    save_memory()

    # LLM
    reply = ask_llm(user_text)
    MEMORY["history"].append({"assistant": reply})
    MEMORY["history"] = MEMORY["history"][-200:]
    save_memory()

    # TTS
    mp3 = tmpdir / "out.mp3"
    tts(reply, str(mp3), tts_lang or "en")

    # HTML with canvas + JS lip sync
    html = f"""
    <link rel="stylesheet" href="style.css">
    <div id="avatar-wrapper">
        <div id="avatar-container">
            <canvas id="avatar-canvas"></canvas>
        </div>
        <audio id="avatar-audio" src="{mp3}" controls></audio>
    </div>

    <script src="avatar.js"></script>
    <script>
        setTimeout(() => {{
            const audio = document.getElementById("avatar-audio");
            audio.oncanplay = () => {{
                audio.play().catch(()=>{});
                if (window.startLipSync) {{
                    window.startLipSync(audio.src);
                }}
            }};
        }}, 300);
    </script>
    """

    return html, reply, str(mp3)

# -----------------------------------------------------------
# ----------------------- GRADIO UI -------------------------
# -----------------------------------------------------------

with gr.Blocks(title="🌸 Animated Avatar Companion") as demo:
    gr.Markdown("# 🌸 Animated Avatar Companion")
    gr.Markdown("Speak or type. The avatar listens, thinks, talks, and animates.")

    with gr.Row():
        with gr.Column(scale=1):
            avatar_upload = gr.File(label="Upload avatar (optional)")
            tts_lang = gr.Textbox(label="TTS language", value="en")
            typed = gr.Textbox(label="Type message (optional)")
            mic = gr.Audio(label="Microphone", sources=["microphone"])
            btn = gr.Button("Send")

        with gr.Column(scale=1):
            html_out = gr.HTML()
            reply_out = gr.Textbox(label="Reply")
            audio_out = gr.Audio(label="Voice", type="file")

    btn.click(
        process,
        inputs=[avatar_upload, mic, typed, tts_lang],
        outputs=[html_out, reply_out, audio_out],
    )

demo.launch()