"""Threaded Fish Speech inference queue for the canonical MXFP8 checkpoint."""

from __future__ import annotations

import queue
import threading
import traceback
from pathlib import Path

import torch
from loguru import logger

from fish_speech.models.text2semantic.inference import (
    GenerateRequest,
    WrappedGenerateResponse,
    decode_one_token_ar,
    generate_long,
)

from .modules import MXFP8Linear, load_mxfp8_checkpoint


def launch_mxfp8_thread_safe_queue(
    checkpoint_path: str | Path,
    device: str | torch.device,
    precision: torch.dtype,
    compile: bool = False,
    *,
    max_length: int = 4096,
    verify_checksums: bool = False,
) -> queue.Queue:
    """Load MXFP8 once and serialize Fish Speech generation through a queue."""

    if precision is not torch.bfloat16:
        raise ValueError("The qualified MXFP8 service requires BF16 exclusions")
    if compile:
        raise ValueError("torch.compile is not qualified for the MXFP8 API service")
    if max_length < 1024:
        raise ValueError("MXFP8 API cache length must be at least 1024")

    input_queue: queue.Queue = queue.Queue()
    init_event = threading.Event()
    init_error: list[BaseException] = []

    def worker() -> None:
        try:
            model = load_mxfp8_checkpoint(
                checkpoint_path,
                device=device,
                max_length=max_length,
                verify_checksums=verify_checksums,
            )
            with torch.device(device):
                model.setup_caches(
                    max_batch_size=1,
                    max_seq_len=model.config.max_seq_len,
                    dtype=torch.bfloat16,
                )
            model._cache_setup_done = True
            module_count = sum(
                isinstance(module, MXFP8Linear) for module in model.modules()
            )
            if module_count != 180:
                raise RuntimeError(
                    f"Expected 180 native MXFP8 projections, found {module_count}"
                )
            logger.info(
                "Loaded qualified MXFP8 slow-transformer checkpoint with "
                f"{module_count} native projections on {device}"
            )
        except BaseException as error:
            init_error.append(error)
            logger.error(traceback.format_exc())
            init_event.set()
            return

        init_event.set()
        while True:
            item: GenerateRequest | None = input_queue.get()
            if item is None:
                break
            response_queue = item.response_queue
            try:
                for chunk in generate_long(
                    model=model,
                    decode_one_token=decode_one_token_ar,
                    **item.request,
                ):
                    response_queue.put(
                        WrappedGenerateResponse(status="success", response=chunk)
                    )
            except Exception as error:
                logger.error(traceback.format_exc())
                response_queue.put(
                    WrappedGenerateResponse(status="error", response=error)
                )
            finally:
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

    threading.Thread(
        target=worker,
        name="fish-s2-pro-mxfp8-worker",
        daemon=True,
    ).start()
    init_event.wait()
    if init_error:
        raise RuntimeError("MXFP8 model worker failed to initialize") from init_error[0]
    return input_queue