from __future__ import annotations
import spaces
import os
import random
import re
import tempfile
import traceback
import requests
from io import BytesIO
from pathlib import Path
from typing import Any, Optional

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import torch

torch._dynamo.config.suppress_errors = True
torch._dynamo.reset()
import whisper
from transformers import AutoModelForCausalLM, AutoTokenizer
from voxcpm import VoxCPM
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from gradio import Server

from db import (
    get_latest_profile_audio,
    list_profiles_for_user,
    list_relationships_for_user,
    relationship_exists_for_user,
    save_profile,
    save_story_to_qdrant,
    search_qdrant,
    debug_database,
)
from loguru import logger

# ── Startup DB check ─────────────────────────────────────────────────────────
try:
    debug_database()
except Exception:
    logger.error(f"Startup DB debug failed: {traceback.format_exc()}")



# ── Lazy model cache (local models only: Whisper + VoxCPM2) ───────────────────
_voxcpm_model = None
_transcriber = None
# ── Lazy model cache ──────────────────────────────────────────────────────────
_story_gen_model = None
_story_gen_tokenizer = None

TRANSCRIBER_MODEL = "base"

RELATION_TRAITS = {
    "grandmother": ["warm", "gentle", "loving", "nostalgic", "tender"],
    "grandfather": ["kind", "steady", "wise", "calm", "patient"],
    "mother":      ["tender", "cheerful", "caring", "warm", "loving"],
    "father":      ["kind", "patient", "gentle", "steady", "calm"],
    "sibling":     ["cheerful", "playful", "kind", "energetic"],
    "friend":      ["cheerful", "warm", "kind", "enthusiastic"],
}
DEFAULT_TRAITS = ["kind", "warm", "loving", "calm", "cheerful", "gentle"]


# ── Local model loaders (Whisper + VoxCPM2 only) ─────────────────────────────

def get_story_gen_model():
    global _story_gen_model, _story_gen_tokenizer
    if _story_gen_model is None:
        model_id = "Qwen/Qwen2.5-3B-Instruct"
        logger.info(f"Loading Story Generation model: {model_id}...")
        try:
            _story_gen_tokenizer = AutoTokenizer.from_pretrained(
                model_id,
                trust_remote_code=True,
            )
            _story_gen_model = AutoModelForCausalLM.from_pretrained(
                model_id,
                trust_remote_code=True,
                torch_dtype=torch.bfloat16,
            )
            _story_gen_model.eval()
            logger.info("Qwen2.5-3B-Instruct loaded successfully on CPU")
        except Exception:
            logger.error(f"Failed to load story gen model: {traceback.format_exc()}")
            raise
    return _story_gen_model, _story_gen_tokenizer

def _generate_story_text_qwen(model, tokenizer, prompt: str) -> str:
    messages = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.8,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True,
    ).strip()
    logger.info(f"Generated text length: {len(response)}")
    return response

@torch._dynamo.disable
def _load_voxcpm():
    """Helper to load VoxCPM with dynamo disabled (must be decorator, not context manager)."""
    return VoxCPM.from_pretrained("openbmb/VoxCPM2")


def get_voxcpm_model():
    global _voxcpm_model
    if _voxcpm_model is None:
        logger.info("Loading VoxCPM2 voice synthesis model...")
        try:
            _voxcpm_model = _load_voxcpm()
            logger.info("VoxCPM2 model loaded successfully (CPU).")
        except Exception:
            logger.error(f"Failed to load VoxCPM2 model: {traceback.format_exc()}")
            raise
    return _voxcpm_model


def get_transcriber():
    global _transcriber
    if _transcriber is None:
        logger.info(f"Loading Whisper model: {TRANSCRIBER_MODEL}...")
        _transcriber = whisper.load_model(TRANSCRIBER_MODEL, device="cpu")
        logger.info("Whisper model loaded successfully.")
    return _transcriber


def preload_local_models():
    """Download and cache Whisper + VoxCPM2 at startup."""
    logger.info("=" * 50)
    logger.info("PRELOADING LOCAL MODELS AT STARTUP")
    logger.info("=" * 50)
    try:
        logger.info("Preloading Whisper...")
        get_transcriber()
        logger.info("Preloading Qwen2.5-3B-Instruct...") 
        get_story_gen_model()  
        logger.info("Preloading VoxCPM2...")
        get_voxcpm_model()
        logger.info("=" * 50)
        logger.info("ALL LOCAL MODELS PRELOADED SUCCESSFULLY")
        logger.info("=" * 50)
    except Exception:
        logger.error(f"Model preloading failed: {traceback.format_exc()}")
        raise



# ── Audio utilities ───────────────────────────────────────────────────────────

def _extract_path_from_audio_dict(audio: dict) -> str | None:
    for key in ("path", "name", "url", "tmp_path"):
        val = audio.get(key)
        if val and isinstance(val, str) and (os.path.exists(val) or val.startswith("http")):
            return val
    for val in audio.values():
        if isinstance(val, str) and len(val) > 1:
            return val
    return None


def validate_audio_input(audio: Any, label: str = "audio") -> dict[str, Any]:
    logger.info(f"--- Validating {label} ---")
    diagnostics: dict[str, Any] = {
        "label": label,
        "type": str(type(audio)),
        "is_none": audio is None,
        "valid": False,
        "error": None,
    }

    if audio is None:
        diagnostics["error"] = "Audio is None"
        logger.warning(f"{label} is None")
        return diagnostics

    if isinstance(audio, dict):
        logger.info(f"{label} received as dict: {list(audio.keys())}")
        path = _extract_path_from_audio_dict(audio)
        if not path:
            diagnostics["error"] = f"Audio dict has no usable path: {audio}"
            logger.error(diagnostics["error"])
            return diagnostics
        diagnostics["type"] = "dict→file"
        audio = path

    if isinstance(audio, (str, Path)):
        path = str(audio)
        diagnostics["path"] = path
        diagnostics["exists"] = os.path.exists(path)
        if diagnostics["exists"]:
            diagnostics["size"] = os.path.getsize(path)
            logger.info(f"{label} is a file: {path} ({diagnostics['size']} bytes)")
            diagnostics["valid"] = diagnostics["size"] > 0
            if not diagnostics["valid"]:
                diagnostics["error"] = "File is empty"
        else:
            diagnostics["error"] = "File does not exist"
            logger.error(f"{label} file not found: {path}")
        return diagnostics

    if isinstance(audio, tuple):
        diagnostics["is_tuple"] = True
        diagnostics["tuple_len"] = len(audio)
        if len(audio) == 2:
            sr, data = audio
            diagnostics["sample_rate"] = sr
            diagnostics["data_type"] = str(type(data))
            if isinstance(data, np.ndarray):
                diagnostics["shape"] = data.shape
                diagnostics["size"] = data.size
                diagnostics["dtype"] = str(data.dtype)
                logger.info(f"{label} is numpy: SR={sr}, shape={data.shape}, dtype={data.dtype}")
                diagnostics["valid"] = data.size > 0
                if not diagnostics["valid"]:
                    diagnostics["error"] = "Numpy array is empty"
            else:
                diagnostics["error"] = "Second element of tuple is not a numpy array"
        else:
            diagnostics["error"] = f"Expected tuple of length 2, got {len(audio)}"
        return diagnostics

    diagnostics["error"] = f"Unsupported audio type: {type(audio)}"
    logger.warning(f"Unknown audio type for {label}: {type(audio)}")
    return diagnostics


def normalize_audio_input(audio: Any) -> tuple[np.ndarray, int]:
    if audio is None:
        raise ValueError("No audio was provided.")

    if isinstance(audio, dict):
        logger.info(f"Audio is a dict with keys: {list(audio.keys())}")
        path = _extract_path_from_audio_dict(audio)
        if not path:
            raise ValueError(f"Audio dict has no usable path: {audio}")
        logger.info(f"Extracted path from dict: {path}")
        audio = path

    if isinstance(audio, (str, Path)):
        logger.info(f"Loading audio from file using librosa: {audio}")
        data, samplerate = librosa.load(str(audio), sr=None, mono=True)
        return data, int(samplerate)

    if isinstance(audio, tuple):
        samplerate, data = audio
        if isinstance(data, np.ndarray):
            if data.dtype == np.int16:
                data = data.astype(np.float32) / 32768.0
            elif data.dtype == np.int32:
                data = data.astype(np.float32) / 2147483648.0
            if len(data.shape) > 1:
                data = np.mean(data, axis=1)
            return data, int(samplerate)

    raise ValueError(f"Unsupported audio type: {type(audio)}")


def audio_to_bytes(audio: Any) -> bytes:
    data, sample_rate = normalize_audio_input(audio)
    buffer = BytesIO()
    sf.write(buffer, data, sample_rate, format="WAV")
    return buffer.getvalue()


def write_temp_audio(audio_bytes: bytes) -> str:
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp.write(audio_bytes)
    tmp.flush()
    tmp.close()
    return tmp.name


# ── Core business logic ───────────────────────────────────────────────────────

def transcribe_audio_story(audio: Any) -> str:
    if isinstance(audio, dict):
        path = _extract_path_from_audio_dict(audio)
        if path:
            audio = path

    diag = validate_audio_input(audio, "story_audio")
    if not diag["valid"]:
        logger.warning(f"Invalid story audio: {diag['error']}")
        return ""

    if isinstance(audio, (str, Path)):
        audio_path = str(audio)
    else:
        audio_path = write_temp_audio(audio_to_bytes(audio))

    logger.info(f"Transcribing audio story from {audio_path}...")
    try:
        model = get_transcriber()
        result = model.transcribe(audio_path)
        text = result.get("text", "").strip()
        logger.info(f"Transcription complete: {text[:50]}...")
        return text
    except Exception as exc:
        logger.error(f"Transcription failed: {traceback.format_exc()}")
        return f"[transcription error: {exc}]"


def build_profile_table(hf_username: str) -> list[list[str]]:
    try:
        items = list_profiles_for_user(hf_username)
        return [[item["id"], item["relationship"], item["created_at"]] for item in items]
    except Exception:
        logger.error(f"Error building profile table: {traceback.format_exc()}")
        return []


@torch._dynamo.disable
def _run_voxcpm_generate(vox, text, reference_wav_path, cfg_value, inference_timesteps):
    """Helper to run VoxCPM generation with dynamo disabled."""
    return vox.generate(
        text=text,
        reference_wav_path=reference_wav_path,
        prompt_wav_path=reference_wav_path,
        inference_timesteps=20,
        cfg_value=2.0,
        prompt_text="The warmth of family is the greatest gift we share. Every laugh, every tear, every quiet moment together becomes a treasure in our hearts. Time may pass, but these memories remain forever, glowing softly like stars in the night sky, guiding us home."
    )

def _generate_story_text_qwen(model, tokenizer, prompt: str) -> str:
    messages = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.8,
            top_p=0.9,
            do_sample=True,
            repetition_penalty=1.1,
        )

    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True,
    ).strip()
    logger.info(f"Generated text length: {len(response)}")
    return response
# ── SINGLE GPU FUNCTION: entire pipeline in one ZeroGPU call ──────────────────

@spaces.GPU
def generate_and_speak(hf_username: str, relationship: str, question: str = "") -> dict:
    """
    Single ZeroGPU function that handles the full pipeline.
    Text generation is done via external MiniCPM4.1-8B API (CPU/network).
    Voice synthesis is done via local VoxCPM2 (GPU).
    """
    logger.info(f"Listen Mode: Generating response for {hf_username}/{relationship} | Question: {question}")
    if not hf_username or not relationship:
        return {"success": False, "status": "Please log in and select a relationship.", "audio_url": None, "text": ""}

    audio_bytes = get_latest_profile_audio(hf_username, relationship)
    if audio_bytes is None:
        logger.warning(f"Reference audio not found for {hf_username}/{relationship}")
        return {"success": False, "status": "No reference voice profile found.", "audio_url": None, "text": ""}

    ref_audio_path = write_temp_audio(audio_bytes)

    try:
        # --- Step 0: Context retrieval (CPU) ---
        context_memory = search_qdrant(hf_username, relationship, question) if question.strip() else None

        if context_memory:
            logger.info("Using Qdrant context for generation.")
            prompt = (
                f"You are a warm, loving {relationship}. "
                f"Your family member asked: '{question}'\n\n"
                f"Draw from these memories: {context_memory}\n\n"
                f"Respond with genuine emotion, specific details, and warmth. "
                f"Speak directly to them. Keep it to 1-2 short paragraphs.you are not allowed to exceed 75 words"
            )
        else:
            logger.info("No context found. Using generic heartwarming prompt.")
            prompt = (
                f"You are a warm, loving {relationship}. "
                f"Your family member asked: '{question}'\n\n"
                f"Make up a story which is only of 2-3 short paragraphs which answers the question"
                f"Make it feel real, emotional, and specific. "
                f"Speak directly to them."
            )

        # --- Step 1: Text Generation (GPU) ---
        logger.info("GPU: Loading Qwen2.5-3B into VRAM...")
        model, tokenizer = get_story_gen_model()
        model = model.cuda()
        story_text = _generate_story_text_qwen(model, tokenizer, prompt)

        # Offload LLM to free VRAM for TTS
        model = model.cpu()
        del model
        torch.cuda.empty_cache()
        logger.info("GPU: Qwen2.5-3B offloaded, VRAM freed.")

        # --- Step 1.5: Post-process text (CPU) ---
        if not context_memory and question.strip():
            story_text = f"My memory fades but : {story_text}"

        relation_lower = relationship.lower()
        trait_pool = DEFAULT_TRAITS
        for key, traits in RELATION_TRAITS.items():
            if key in relation_lower:
                trait_pool = traits
                break
        chosen_trait = random.choice(trait_pool)
        style_prefix = f"(slow, clear, {chosen_trait} tone) "
        styled_story = story_text

        # --- Step 2: Voice Synthesis (GPU) ---
        logger.info("GPU: Loading VoxCPM2 into VRAM...")
        vox = get_voxcpm_model()

        # Move internal modules to CUDA
        if hasattr(vox, "to"):
            vox = vox.to("cuda")
        else:
            for attr_name in ["model", "tts_model", "vae", "encoder", "decoder"]:
                if hasattr(vox, attr_name):
                    attr = getattr(vox, attr_name)
                    if hasattr(attr, "cuda"):
                        setattr(vox, attr_name, attr.cuda())

        output_audio_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name

        wav = _run_voxcpm_generate(
            vox,
            text=styled_story,
            reference_wav_path=ref_audio_path,
            cfg_value=3.0,
            inference_timesteps=10,
        )
        sf.write(output_audio_path, wav, vox.tts_model.sample_rate)

        # Cleanup TTS from VRAM
        if hasattr(vox, "cpu"):
            vox = vox.cpu()
        else:
            for attr_name in ["model", "tts_model", "vae", "encoder", "decoder"]:
                if hasattr(vox, attr_name):
                    attr = getattr(vox, attr_name)
                    if hasattr(attr, "cpu"):
                        setattr(vox, attr_name, attr.cpu())
        del vox
        torch.cuda.empty_cache()
        logger.info("GPU: VoxCPM2 offloaded, VRAM freed.")

        filename = os.path.basename(output_audio_path)
        return {
            "success": True,
            "status": f"Generated for {relationship} ({chosen_trait} tone).",
            "audio_url": f"/audio/{filename}",
            "text": story_text,
        }
    except Exception as exc:
        logger.error(f"Generation/Synthesis failed: {traceback.format_exc()}")
        return {"success": False, "status": f"Generation error: {exc}", "audio_url": None, "text": ""}


# ── API handler functions ─────────────────────────────────────────────────────

def handle_login(hf_username: str) -> dict:
    username = (hf_username or "").strip()
    logger.info(f"Login requested for username: {username}")
    if not username:
        logger.warning("Login attempted with empty username.")
        return {
            "success": False,
            "status": "Please enter your Hugging Face username.",
            "username": "",
            "profiles": [],
            "relationships": [],
        }
    profiles = build_profile_table(username)
    relationships = list_relationships_for_user(username)
    logger.info(f"Login successful for {username}. Profiles: {len(profiles)}. Relationships: {len(relationships)}.")
    return {
        "success": True,
        "status": f"✅ Welcome, {username}!",
        "username": username,
        "profiles": profiles,
        "relationships": relationships,
    }


def save_voice_profile_only(hf_username: str, relationship: str, audio: Any) -> dict:
    logger.info("=== SAVE PROFILE START ===")
    logger.info(f"User: {hf_username} | Relationship: {relationship}")
    relation = (relationship or "").strip()

    if isinstance(audio, dict):
        path = _extract_path_from_audio_dict(audio)
        if path:
            logger.info(f"Unwrapped audio dict → {path}")
            audio = path

    audio_diag = validate_audio_input(audio, "profile_audio")

    if not hf_username:
        return {"success": False, "status": "Please log in first.", "profiles": [], "relationships": []}
    if not relation:
        return {"success": False, "status": "Enter a relationship label.", "profiles": build_profile_table(hf_username), "relationships": list_relationships_for_user(hf_username)}
    if not audio_diag["valid"]:
        return {"success": False, "status": f"Voice sample issue: {audio_diag['error'] or 'Unknown error'}", "profiles": build_profile_table(hf_username), "relationships": list_relationships_for_user(hf_username)}

    try:
        audio_bytes = audio_to_bytes(audio)
        logger.info(f"Converting profile audio: {len(audio_bytes)} bytes")
        save_profile(hf_username, relation, audio_bytes)
        logger.info("MongoDB save successful.")
        profiles = build_profile_table(hf_username)
        relationships = list_relationships_for_user(hf_username)
        logger.info("=== SAVE PROFILE COMPLETED SUCCESSFULLY ===")
        return {"success": True, "status": "✅ Voice profile saved successfully.", "profiles": profiles, "relationships": relationships}
    except Exception as exc:
        logger.error(f"Failed to save voice profile: {traceback.format_exc()}")
        return {"success": False, "status": f"Error: {exc}", "profiles": build_profile_table(hf_username), "relationships": list_relationships_for_user(hf_username)}


def save_story_memory(hf_username: str, relationship: str, story_audio: Any) -> dict:
    logger.info("=== SAVE STORY START ===")
    logger.info(f"User: {hf_username} | Relationship: {relationship}")
    relation = (relationship or "").strip()

    if isinstance(story_audio, dict):
        path = _extract_path_from_audio_dict(story_audio)
        if path:
            logger.info(f"Unwrapped story_audio dict → {path}")
            story_audio = path

    story_diag = validate_audio_input(story_audio, "story_audio")

    if not hf_username:
        return {"success": False, "status": "Please log in first.", "transcript": ""}
    if not relation:
        return {"success": False, "status": "Please select a relationship for this memory.", "transcript": ""}
    if not story_diag["valid"]:
        return {"success": False, "status": f"Story audio issue: {story_diag['error'] or 'Unknown error'}", "transcript": ""}

    try:
        logger.info("Starting Whisper transcription...")
        transcript = transcribe_audio_story(story_audio)
        logger.info(f"Transcription complete (length: {len(transcript)})")
        if not transcript:
            return {"success": False, "status": "Transcription was empty. Try speaking more clearly.", "transcript": ""}
        logger.info(f"Saving to Qdrant collection: {hf_username}-{relation}")
        save_story_to_qdrant(hf_username, relation, transcript)
        logger.info("Qdrant storage successful.")
        logger.info("=== SAVE STORY COMPLETED SUCCESSFULLY ===")
        return {"success": True, "status": "✅ Memory saved successfully.", "transcript": transcript}
    except Exception as exc:
        logger.error(f"Failed to save story memory: {traceback.format_exc()}")
        return {"success": False, "status": f"Error: {exc}", "transcript": ""}


def refresh_profiles(hf_username: str) -> dict:
    relationships = list_relationships_for_user(hf_username)
    profiles = build_profile_table(hf_username)
    return {"profiles": profiles, "relationships": relationships}


def load_voice_profile(hf_username: str, relationship: str) -> dict:
    if not hf_username:
        return {"success": False, "status": "Please log in first.", "audio_url": None}
    if not relationship:
        return {"success": False, "status": "Choose a relationship from the list.", "audio_url": None}
    audio_bytes = get_latest_profile_audio(hf_username, relationship)
    if audio_bytes is None:
        return {"success": False, "status": "No saved voice profile found for that relationship.", "audio_url": None}
    audio_path = write_temp_audio(audio_bytes)
    filename = os.path.basename(audio_path)
    return {"success": True, "status": f"Playing saved voice for {relationship}.", "audio_url": f"/audio/{filename}"}


# ── gradio.Server app ─────────────────────────────────────────────────────────

FRONTEND_DIR = Path(__file__).parent / "static"

app = Server()


@app.api(name="login")
def api_login(hf_username: str) -> dict:
    return handle_login(hf_username)


@app.api(name="save_profile")
def api_save_profile(hf_username: str, relationship: str, audio: Optional[str]) -> dict:
    return save_voice_profile_only(hf_username, relationship, audio)


@app.api(name="save_story")
def api_save_story(hf_username: str, relationship: str, audio: Optional[str]) -> dict:
    return save_story_memory(hf_username, relationship, audio)


@app.api(name="refresh_profiles")
def api_refresh(hf_username: str) -> dict:
    return refresh_profiles(hf_username)


@app.api(name="load_voice")
def api_load_voice(hf_username: str, relationship: str) -> dict:
    return load_voice_profile(hf_username, relationship)


@app.api(name="generate")
def api_generate(hf_username: str, relationship: str, question: str) -> dict:
    return generate_and_speak(hf_username, relationship, question)


@app.api(name="debug_audio")
def api_debug_audio(audio: Optional[str], story_audio: Optional[str]) -> str:
    diag1 = validate_audio_input(audio, "profile_audio")
    diag2 = validate_audio_input(story_audio, "story_audio")
    return f"Profile Audio: {diag1}\n\nStory Audio: {diag2}"


# ── Serve generated audio files ──────────────────────────────────────────────

@app.get("/audio/{filename}")
async def serve_audio(filename: str):
    if not filename.endswith(".wav") or ".." in filename or "/" in filename:
        return {"error": "Invalid filename"}

    file_path = os.path.join("/tmp", filename)
    if os.path.exists(file_path):
        return FileResponse(file_path, media_type="audio/wav")
    return {"error": "File not found"}


# ── Serve custom frontend ─────────────────────────────────────────────────────

@app.get("/", response_class=HTMLResponse)
async def homepage() -> HTMLResponse:
    index_path = FRONTEND_DIR / "index.html"
    return HTMLResponse(index_path.read_text(encoding="utf-8"))


if FRONTEND_DIR.exists():
    app.mount("/static", StaticFiles(directory=str(FRONTEND_DIR)), name="static")


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    # Preload Whisper + VoxCPM2 only. MiniCPM4.1-8B is external API.
    preload_local_models()
    app.launch(show_error=True, server_port=7860)