# app.py - Fixed version (Microphone-only input)
# Drop this into your HF Space alongside: avatar.png, avatar.js, style.css
# Requirements: gradio, transformers, torch, gtts, pydub, faster-whisper (optional)
# Notes: This version uses microphone-only input (gr.Audio sources=["microphone"])

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

import gradio as gr

# STT (prefer faster_whisper if available; fallback to whisper)
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: use a lightweight ungated model that's safe for public Spaces
# Flan-T5 small is reliable on CPU and avoids gated repos
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 Exception as e:
    llm_pipe = None
    print("LLM load error:", e)

# TTS
from gtts import gTTS
from pydub import AudioSegment

# Memory persistence
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 Exception:
        MEMORY = {"history": []}
else:
    MEMORY = {"history": []}


def save_memory():
    try:
        with open(MEMORY_FILE, "w", encoding="utf-8") as f:
            json.dump(MEMORY, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print("Memory save error:", e)


# ---------------- Helpers ----------------
def transcribe_audio(path):
    """Transcribe audio using faster_whisper or 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 Exception as e:
        print("Transcription error:", e)
        return ""


def ask_llm(user_text, max_length=200):
    """Ask the LLM and return a concise reply with memory context."""
    if llm_pipe is None:
        return "Sorry — language model temporarily unavailable."
    try:
        # Build a compact context from memory
        recent = MEMORY.get("history", [])[-8:]
        ctx_lines = []
        for item in recent:
            if "user" in item:
                ctx_lines.append(f"User: {item['user']}")
            if "assistant" in item:
                ctx_lines.append(f"Assistant: {item['assistant']}")
        ctx = "\n".join(ctx_lines)

        prompt = (
            "You are a kind, highly capable AI companion and tutor. Answer concisely, "
            "give a short practical example when relevant, and ask one follow-up question. "
            "Match the user's language and tone.\n\n"
            f"{ctx}\nUser: {user_text}\nAssistant:"
        )
        out = llm_pipe(prompt, max_length=max_length, do_sample=False)
        text = out[0]["generated_text"]
        # Trim possible preamble
        if "Assistant:" in text:
            text = text.split("Assistant:")[-1].strip()
        return text.strip()
    except Exception as e:
        print("LLM error:", e)
        return "I couldn't think clearly — try asking again."


def tts_gtts_save(text, out_path, lang="en"):
    """Create an mp3 via gTTS (multilingual)."""
    try:
        if not text or not text.strip():
            silent = AudioSegment.silent(duration=400)
            silent.export(out_path, format="mp3")
            return out_path
        tts = gTTS(text=text, lang=lang, slow=False)
        tts.save(out_path)
        return out_path
    except Exception as e:
        print("TTS error:", e)
        silent = AudioSegment.silent(duration=400)
        silent.export(out_path, format="mp3")
        return out_path


# ---------------- Main pipeline ----------------
def process_interaction(avatar_upload, mic_audio, typed_text, tts_lang):
    """
    microphone-only mode: mic_audio will be provided by gradio
    avatar_upload: optional file upload (if user wants to override repo avatar)
    typed_text: optional typed message
    tts_lang: language code for TTS
    """
    uid = uuid.uuid4().hex[:8]
    tmpdir = Path(tempfile.gettempdir()) / f"avatar_{uid}"
    tmpdir.mkdir(parents=True, exist_ok=True)

    # Avatar selection: prefer repo avatar.png if exists, else avatar_upload
    repo_avatar = Path("avatar.png")
    if repo_avatar.exists():
        avatar_path = str(repo_avatar)
    elif avatar_upload:
        # avatar_upload is a file-like object (UploadedFile); save it
        try:
            p = tmpdir / "avatar_upload.png"
            with open(p, "wb") as f:
                f.write(avatar_upload.read())
            avatar_path = str(p)
        except Exception:
            avatar_path = None
    else:
        avatar_path = None

    if avatar_path is None:
        return gr.HTML("<b>No avatar found.</b> Upload avatar.png to the repo or use the upload field."), "No avatar", None

    # Determine user text
    user_text = ""
    if typed_text and typed_text.strip():
        user_text = typed_text.strip()
    else:
        if mic_audio is None:
            return None, "Provide audio (microphone) or type a message.", None

        # mic_audio is a tempfile-like object with .read(); save it
        audio_path = tmpdir / "user_in.wav"
        try:
            with open(audio_path, "wb") as f:
                f.write(mic_audio.read())
        except Exception as e:
            print("Audio save error:", e)
            return None, "Failed to read microphone audio.", None

        user_text = transcribe_audio(str(audio_path))

    if not user_text:
        return None, "Couldn't transcribe. Try again or type your message.", None

    # Add to memory
    MEMORY.setdefault("history", []).append({"user": user_text, "time": datetime.utcnow().isoformat()})
    MEMORY["history"] = MEMORY["history"][-200:]
    save_memory()

    # LLM reply
    reply_text = ask_llm(user_text)

    # Save assistant reply to memory
    MEMORY.setdefault("history", []).append({"assistant": reply_text, "time": datetime.utcnow().isoformat()})
    MEMORY["history"] = MEMORY["history"][-200:]
    save_memory()

    # TTS
    tts_path = tmpdir / "reply.mp3"
    tts_gtts_save(reply_text, str(tts_path), lang=tts_lang or "en")

    # Build HTML to show avatar canvas and trigger lip-sync using avatar.js
    # Use Gradio-served file path (client will resolve)
    audio_file_url = str(tts_path)
    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="{audio_file_url}" controls></audio>
    </div>
    <script src="avatar.js"></script>
    <script>
      setTimeout(function() {{
        const a = document.getElementById('avatar-audio');
        a.oncanplay = () => {{
          a.play().catch(()=>{{}}); 
          if (window.startLipSync) {{
              try {{ window.startLipSync(a.src); }} catch(e){{console.log(e);}}
          }}
        }};
        if (a.readyState >= 2) {{
          a.play().catch(()=>{{}});
          if (window.startLipSync) {{
              try {{ window.startLipSync(a.src); }} catch(e){{console.log(e);}}
          }}
        }}
      }}, 500);
    </script>
    """

    return html, reply_text, str(tts_path)


# ---------------- Gradio UI ----------------
title = "🌸 Live Animated Avatar Companion — Mic only"
description = (
    "Place avatar.png in the Space repo (preferred). Use the mic to speak (microphone-only mode). "
    "Or optionally upload an avatar image (override). The avatar will reply with text + voice and animate in-browser "
    "(lip-sync, blinking, head motion) using avatar.js."
)

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)

    with gr.Row():
        with gr.Column(scale=1):
            avatar_upload = gr.File(label="Upload avatar image (optional) - repo avatar.png preferred")
            tts_lang = gr.Textbox(value="en", label="TTS language (e.g., en, es, fr)", max_lines=1)
            typed_text = gr.Textbox(lines=2, placeholder="Or type your message (optional)", label="Type message")
            # Microphone-only audio input (new gradio API)
            mic_audio = gr.Audio(label="Record via Microphone (required for voice)", sources=["microphone"])
            submit = gr.Button("Send")

        with gr.Column(scale=1):
            html_out = gr.HTML(label="Animated avatar (plays reply)")
            txt_out = gr.Textbox(label="Assistant reply")
            # fixed: use 'filepath' which is a valid Gradio audio output type
            audio_out = gr.Audio(label="Reply audio", type="filepath")

    def on_submit(avatar_upload_, mic_audio_, typed_text_, tts_lang_):
        return process_interaction(avatar_upload_, mic_audio_, typed_text_, tts_lang_)

    submit.click(fn=on_submit,
                 inputs=[avatar_upload, mic_audio, typed_text, tts_lang],
                 outputs=[html_out, txt_out, audio_out])

    # memory controls
    with gr.Row():
        view_btn = gr.Button("View recent memory")
        clear_btn = gr.Button("Clear memory")
        mem_display = gr.Textbox(label="Memory (last 20)")

    def view_memory():
        return json.dumps(MEMORY.get("history", [])[-20:], ensure_ascii=False, indent=2)

    def clear_memory():
        MEMORY["history"] = []
        save_memory()
        return "Memory cleared."

    view_btn.click(view_memory, inputs=None, outputs=mem_display)
    clear_btn.click(clear_memory, inputs=None, outputs=mem_display)

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