"""Kith Local — Build Small spike.

A relationship memory on small open models you own. Audio → Cohere Transcribe
(2B, on ZeroGPU) → MiniCPM extraction (OpenBMB) → SQLite → the Today feed,
served as a hand-written HTML page over gr.Server().
"""

import base64
import os
import tempfile
from datetime import date

import gradio as gr
import spaces
import torch
from fastapi.responses import FileResponse, HTMLResponse

import brain
import db
from frontend import FRONTEND_HTML

ASR_MODEL_ID = "CohereLabs/cohere-transcribe-03-2026"
BRAIN_MODEL_ID = os.environ.get("KITH_GPU_MODEL", "openbmb/MiniCPM4.1-8B")
HF_TOKEN = os.environ.get("HF_TOKEN")

device = "cuda" if torch.cuda.is_available() else "cpu"

# ZeroGPU gotcha: the GPU attaches ONLY inside @spaces.GPU functions, so
# torch.cuda.is_available() is False at module import on a ZeroGPU Space.
# Detect the GPU Space via the env var the `spaces` runtime sets instead.
ON_GPU_SPACE = bool(os.environ.get("SPACES_ZERO_GPU")) or torch.cuda.is_available()

# Dev stub: KITH_SKIP_ASR=1 runs the app without the gated ASR weights
# (extraction/feed/UI still fully live). Never set on the Space.
SKIP_ASR = os.environ.get("KITH_SKIP_ASR") == "1"

# Brain backend: on the GPU Space run MiniCPM 4.1-8B LOCALLY — best quality,
# fully private, no dependency on the temporary hosted endpoint. On a CPU dev
# box, fall back to hosted. brain.py read KITH_BRAIN at ITS import, so assign
# the resolved value directly (env set afterward wouldn't propagate).
brain.KITH_BRAIN = os.environ.get("KITH_BRAIN") or ("gpu" if ON_GPU_SPACE else "hosted")

if not SKIP_ASR:
    from transformers import AutoProcessor, CohereAsrForConditionalGeneration

    processor = AutoProcessor.from_pretrained(ASR_MODEL_ID, token=HF_TOKEN)
    asr_model = CohereAsrForConditionalGeneration.from_pretrained(
        ASR_MODEL_ID, dtype="auto", token=HF_TOKEN
    ).to(device)  # module level: ZeroGPU registers weights for fast restore

# Brain model — loaded at module level (ZeroGPU fast-restore), used only when
# the GPU path is active. Skipped on CPU dev boxes (hosted fallback instead).
if brain.KITH_BRAIN == "gpu":
    from transformers import AutoModelForCausalLM, AutoTokenizer

    brain_tokenizer = AutoTokenizer.from_pretrained(BRAIN_MODEL_ID, trust_remote_code=True, token=HF_TOKEN)
    brain_model = AutoModelForCausalLM.from_pretrained(
        BRAIN_MODEL_ID, dtype="auto", trust_remote_code=True, token=HF_TOKEN,
        low_cpu_mem_usage=True,
    ).to(device)

    @spaces.GPU(duration=120)
    def _brain_chat_gpu(messages: list, temperature: float) -> str:
        """Local MiniCPM4.1-8B generate. Same (messages, temp) -> text contract
        brain.extract() expects; thinking disabled, JSON enforced by the prompt."""
        dev = "cuda" if torch.cuda.is_available() else "cpu"
        brain_model.to(dev)
        inputs = brain_tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            enable_thinking=False,
        ).to(dev)
        out = brain_model.generate(
            **inputs,
            max_new_tokens=2500,
            do_sample=temperature > 0,
            temperature=temperature if temperature > 0 else None,
            pad_token_id=brain_tokenizer.eos_token_id,
        )
        gen = out[0][inputs["input_ids"].shape[1]:]
        return brain_tokenizer.decode(gen, skip_special_tokens=True)

    brain.set_gpu_chat(_brain_chat_gpu)


@spaces.GPU(duration=45)
def _transcribe_file(wav_path: str) -> str:
    import librosa

    # cuda IS available here (inside @spaces.GPU) even though it wasn't at import
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    asr_model.to(dev)
    audio, _ = librosa.load(wav_path, sr=16000, mono=True)
    inputs = processor(
        audio=audio,
        sampling_rate=16000,
        return_tensors="pt",
        language="en",
        punctuation=True,
    ).to(dev)
    inputs["input_features"] = inputs["input_features"].to(asr_model.dtype)
    generated = asr_model.generate(**inputs, max_new_tokens=448)
    chunks = processor.batch_decode(generated, skip_special_tokens=True)
    return " ".join(c.strip() for c in chunks).strip()


def _transcribe_stub(_wav_path: str) -> str:
    sample = os.path.join(os.path.dirname(__file__), "data", "sample-transcript.txt")
    with open(sample) as f:
        return f.read().strip()


# --------------------------------------------------------------------------- #
# Server
# --------------------------------------------------------------------------- #
server = gr.Server()


@server.get("/", response_class=HTMLResponse)
async def homepage() -> str:
    return FRONTEND_HTML


@server.get("/sample-audio")
async def sample_audio() -> FileResponse:
    return FileResponse(
        os.path.join(os.path.dirname(__file__), "data", "sample-conversation.mp3"),
        media_type="audio/mpeg",
    )


# Curated demo scenarios the UI can one-click test. Whitelisted (no path traversal).
DEMO_CLIPS = {"sara-wedding", "daniel-career", "aunt-rose-call"}


@server.get("/clip/{name}")
async def demo_clip(name: str) -> FileResponse:
    from fastapi import HTTPException

    if name not in DEMO_CLIPS:
        raise HTTPException(status_code=404, detail="unknown clip")
    return FileResponse(
        os.path.join(os.path.dirname(__file__), "data", "demo-clips", f"{name}.mp3"),
        media_type="audio/mpeg",
    )


@server.api(name="transcribe")
def transcribe_api(audio_b64: str) -> str:
    """16kHz mono WAV, base64-encoded by the browser -> transcript text."""
    raw = base64.b64decode(audio_b64)
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        f.write(raw)
        path = f.name
    try:
        return _transcribe_stub(path) if SKIP_ASR else _transcribe_file(path)
    finally:
        os.unlink(path)


@server.api(name="ingest")
def ingest_api(transcript: str) -> dict:
    """Transcript -> MiniCPM extraction -> SQLite. Returns counts + extraction."""
    today = date.today().isoformat()
    conversation_id = db.insert_conversation(transcript)
    extraction = brain.extract(transcript, today)
    counts = db.apply_extraction(conversation_id, extraction, today)
    return {"counts": counts, "extraction": extraction}


@server.api(name="today")
def today_api(_ping: str = "") -> list:
    """The Today feed: up to 6 cards (thoughtful / owed / reconnect)."""
    return db.today_feed()


@server.api(name="conversations")
def conversations_api(_ping: str = "") -> list:
    """History for the sidebar: every processed conversation, newest first."""
    return db.list_conversations()


@server.api(name="conversation_detail")
def conversation_detail_api(conversation_id: int) -> dict:
    """One past conversation: its transcript + everything extracted from it."""
    return db.conversation_detail(conversation_id)


@server.api(name="card_detail")
def card_detail_api(card_id: str) -> dict:
    """Everything behind one card: source quote + the rest of that person's memory."""
    return db.card_detail(card_id)


@server.api(name="dismiss")
def dismiss_api(card_id: str) -> dict:
    """Hide one card; keep the underlying memory. Returns the refreshed feed."""
    db.dismiss_card(card_id)
    return {"cards": db.today_feed()}


@server.api(name="reset")
def reset_api(_confirm: str = "") -> dict:
    """Forget everything in the local memory."""
    return {"counts": db.reset_all()}


if __name__ == "__main__":
    server.launch(server_name="0.0.0.0", server_port=7860, show_error=True)