"""
LFM2.5-Audio-1.5B inference module — end-to-end audio Q&A.

Replaces the 3-model pipeline (Whisper ASR + Qwen Q&A + Qwen TTS) with a single
multimodal model that accepts audio/text input and produces audio+text output.
"""
from __future__ import annotations

import logging
import os
import re
import threading
import time

if os.name == "nt":
    os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")  # Windows typically lacks Triton for torch.compile

import torch
import numpy as np

from runtime_config import GPU_INFERENCE_LOCK

logger = logging.getLogger(__name__)

_processor = None
_model = None
_model_lock = threading.Lock()

HF_REPO = "LiquidAI/LFM2.5-Audio-1.5B"
SAMPLE_RATE = 24000  # Mimi codec native rate (confirmed in official demo + library constants)

# Matches special boundary tokens like <|text_end|>, <|audio_start|> that leak into decoded text
_SPECIAL_TOKEN_RE = re.compile(r"<\|[^|>]+\|>")


def _select_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _select_dtype() -> torch.dtype:
    if not torch.cuda.is_available():
        return torch.float32
    cap = torch.cuda.get_device_capability()
    return torch.bfloat16 if cap[0] >= 8 else torch.float16


def _move_module(module, device: torch.device, dtype: torch.dtype):
    if not hasattr(module, "to"):
        return module
    try:
        return module.to(device=device, dtype=dtype)
    except TypeError:
        try:
            return module.to(device)
        except Exception:
            return module
    except Exception:
        return module


def _first_parameter_device(module, fallback: torch.device) -> torch.device:
    try:
        return next(module.parameters()).device
    except Exception:
        try:
            return next(module.buffers()).device
        except Exception:
            return fallback


def _module_device(module, fallback: torch.device) -> torch.device:
    try:
        return next(module.parameters()).device
    except Exception:
        return getattr(module, "device", fallback)


def _assemble_waveform(wav_chunks: list) -> np.ndarray | None:
    """Concatenate Mimi output chunks, normalize to peak 0.9, and return float32 array."""
    if not wav_chunks:
        return None
    try:
        waveform = torch.cat(wav_chunks, dim=-1).float().numpy().squeeze()
    except Exception as exc:
        logger.warning("Waveform assembly failed: %s", exc)
        return None

    if waveform.ndim == 0 or waveform.size == 0:
        return None

    peak = float(np.abs(waveform).max())
    logger.info("Waveform peak amplitude: %.6f  (samples: %d)", peak, waveform.size)

    if peak == 0.0:
        logger.warning("Waveform is all zeros — model generated silence")
        return None

    return waveform * (0.9 / peak)


def get_model_status() -> dict:
    """Return current model load state, device, dtype, and GPU memory."""
    status: dict = {
        "model_loaded": _model is not None,
        "repo": HF_REPO,
        "device": None,
        "dtype": None,
        "gpu_name": None,
        "gpu_memory_used_gb": None,
        "gpu_memory_reserved_gb": None,
    }
    if _model is not None:
        device = _module_device(_model, _select_device())
        status["device"] = str(device)
        try:
            status["dtype"] = str(next(_model.parameters()).dtype).replace("torch.", "")
        except Exception:
            pass
        if device.type == "cuda":
            try:
                status["gpu_name"] = torch.cuda.get_device_name(device)
                status["gpu_memory_used_gb"] = round(torch.cuda.memory_allocated(device) / 1e9, 2)
                status["gpu_memory_reserved_gb"] = round(torch.cuda.memory_reserved(device) / 1e9, 2)
            except Exception:
                pass
    return status


def get_lfm_model():
    """Load LFM2.5-Audio-1.5B model. Cached after first call."""
    global _processor, _model
    if _model is None:
        with _model_lock:
            if _model is None:
                from liquid_audio import LFM2AudioModel, LFM2AudioProcessor

                device = _select_device()
                dtype = _select_dtype()
                logger.info("Loading %s on %s (%s)...", HF_REPO, device, dtype)
                _processor = LFM2AudioProcessor.from_pretrained(HF_REPO, device=device).eval()
                try:
                    _model = LFM2AudioModel.from_pretrained(
                        HF_REPO,
                        dtype=dtype if device.type == "cuda" else torch.float32,
                        device=device,
                    ).eval()
                except TypeError:
                    _model = LFM2AudioModel.from_pretrained(HF_REPO).eval()

                moved_processor = _move_module(_processor, device, dtype)
                if moved_processor is not None:
                    _processor = moved_processor
                moved_model = _move_module(_model, device, dtype)
                if moved_model is not None:
                    _model = moved_model
                _model = _model.eval()
                # Force lazy Mimi construction after the processor is on the target device,
                # and fail early if the streaming decoder cannot run there.
                mimi = _processor.mimi.eval()
                if device.type == "cuda":
                    with torch.no_grad(), mimi.streaming(1):
                        mimi.decode(torch.randint(0, 2048, (1, 8, 1), device=device))
                logger.info("LFM2.5-Audio loaded on %s.", _module_device(_model, device))
    return _processor, _model


def warmup_lfm():
    """Run a dummy generation to warm the model after loading.
    Call once at startup after get_lfm_model() to eliminate first-query latency.
    """
    try:
        from liquid_audio import ChatState

        processor, model = get_lfm_model()
        device = _module_device(model, _select_device())
        dtype = next(model.parameters()).dtype

        with GPU_INFERENCE_LOCK, torch.no_grad():
            chat = ChatState(processor, dtype=dtype)
            chat.new_turn("system")
            chat.add_text("Respond with interleaved text and audio.")
            chat.end_turn()
            chat.new_turn("user")
            chat.add_text("Hi")
            chat.end_turn()
            chat.new_turn("assistant")

            for _ in model.generate_interleaved(
                **chat, max_new_tokens=10, audio_temperature=1.0, audio_top_k=4,
            ):
                pass  # Just trigger warmup, discard output

        logger.info("LFM warmup complete.")
    except Exception as exc:
        logger.warning("LFM warmup failed (non-fatal): %s", exc)


def answer_question_audio(
    question_audio_path: str | None = None,
    question_text: str | None = None,
    story_context: str = "",
    max_new_tokens: int = 150,
) -> tuple[str, np.ndarray | None, int]:
    """
    Answer a question about the story using LFM2.5-Audio end-to-end.

    Accepts either audio input (child's voice) or text input.
    Returns (answer_text, audio_waveform_or_None, sample_rate).
    """
    from liquid_audio import ChatState

    processor, model = get_lfm_model()
    device = _module_device(model, _select_device())
    dtype = next(model.parameters()).dtype

    with GPU_INFERENCE_LOCK, torch.no_grad():
        chat = ChatState(processor, dtype=dtype)

        # System prompt: format requirement + brevity constraint
        chat.new_turn("system")
        chat.add_text(
            "Respond with interleaved text and audio. "
            "Give a short, direct answer in 1-2 sentences. Do not repeat the question."
        )
        chat.end_turn()

        # User turn — story context as text prefix, then audio or text question
        chat.new_turn("user")
        if story_context:
            chat.add_text(
                f"Story context:\n{story_context[:2000]}\n\n"
                "Based only on the story above, answer briefly."
            )
        if question_audio_path:
            import librosa
            wav_np, sr = librosa.load(question_audio_path, sr=16000, mono=True)
            wav = torch.from_numpy(wav_np).unsqueeze(0).to(device)
            chat.add_audio(wav, sr)
        elif question_text:
            chat.add_text(question_text)
        else:
            return "Please ask a question!", None, SAMPLE_RATE
        # Closing constraint so the model sees it immediately before generating
        chat.add_text("Answer in 1-2 sentences only.")
        chat.end_turn()

        chat.new_turn("assistant")
        text_out: list[torch.Tensor] = []
        audio_out: list[torch.Tensor] = []
        text_token_count = 0
        t0 = time.perf_counter()

        for t in model.generate_interleaved(
            **chat,
            max_new_tokens=max_new_tokens,
            audio_temperature=1.0,
            audio_top_k=4,
        ):
            if t.numel() == 1:
                text_token_count += 1
                text_out.append(t)
            elif t.numel() == 8:
                audio_out.append(t)

        gen_time = time.perf_counter() - t0
        audio_frame_count = max(0, len(audio_out) - 1)  # last frame is EOS
        logger.info(
            "Generated: %d text tokens, %d audio frames (%.1f sec), %.1f s wall",
            text_token_count, audio_frame_count, audio_frame_count / 12.5, gen_time,
        )

    # Decode text — strip interleaved boundary special tokens that leak through
    answer_text = ""
    if text_out:
        raw_text = "".join(processor.text.decode(t.detach().cpu()) for t in text_out)
        answer_text = _SPECIAL_TOKEN_RE.sub("", raw_text).strip()

    # Decode audio via processor.decode (drops EOS frame)
    waveform = None
    if len(audio_out) > 1:
        audio_codes = torch.stack(audio_out[:-1], dim=1).unsqueeze(0)  # (1, 8, N)
        raw = processor.decode(audio_codes).cpu().float()
        waveform = raw[0].numpy()
        peak = float(np.abs(waveform).max())
        if peak > 0:
            waveform = waveform * (0.9 / peak)
        logger.info("Decoded audio: %.2f sec, peak %.4f", len(waveform) / SAMPLE_RATE, peak)

    return answer_text, waveform, SAMPLE_RATE


def text_to_audio_lfm(
    text: str,
    max_new_tokens: int = 1024,
) -> tuple[np.ndarray | None, int]:
    """Convert text to audio using LFM2.5-Audio in TTS mode.

    Feeds the text as a user message with a "read aloud" system prompt,
    then collects only the audio tokens from generate_interleaved.
    Returns (waveform_or_None, sample_rate).
    """
    from liquid_audio import ChatState

    if not text.strip():
        return None, SAMPLE_RATE

    processor, model = get_lfm_model()
    device = _module_device(model, _select_device())
    dtype = next(model.parameters()).dtype

    with GPU_INFERENCE_LOCK, torch.no_grad():
        chat = ChatState(processor, dtype=dtype)

        chat.new_turn("system")
        chat.add_text("Read the following text aloud clearly and naturally.")
        chat.end_turn()

        chat.new_turn("user")
        chat.add_text(text)
        chat.end_turn()

        chat.new_turn("assistant")
        wav_chunks = []
        audio_frame_count = 0

        mimi = processor.mimi.eval()
        mimi_device = _first_parameter_device(mimi, device)

        with mimi.streaming(1):
            for t in model.generate_interleaved(
                **chat,
                max_new_tokens=max_new_tokens,
                audio_temperature=1.0,
                audio_top_k=4,
            ):
                if t.numel() == 8:
                    if (t == 2048).any():
                        continue
                    audio_frame_count += 1
                    try:
                        wav_chunk = mimi.decode(t[None, :, None].to(device=mimi_device, dtype=torch.long))[0]
                        wav_chunks.append(wav_chunk.cpu())
                    except Exception as exc:
                        logger.warning("TTS decode skipped frame: %s", exc)

        logger.info("TTS: %d audio frames (%.1f sec)", audio_frame_count, audio_frame_count / 12.5)

    return _assemble_waveform(wav_chunks), SAMPLE_RATE