"""Unified inference pipeline for sign language segmentation.

Loads a segmenter v2 checkpoint and runs BIO prediction on WiLoR H5 files.

Adapted from segmenterv2/segmenter/inference/runner.py.
"""

from __future__ import annotations

import logging
import threading
from pathlib import Path

import torch
from torch.nn.utils.rnn import pad_sequence

from sltk.segmentation.h5_loader import h5_to_features, load_h5_directory
from sltk.segmentation.output import OutputFormat, save_predictions
from sltk.segmentation.postprocess import label_fixes, merge_begin_tags

logger = logging.getLogger(__name__)


class InferenceRunner:
    """Runs sign language segmentation inference on H5 feature files.

    Args:
        checkpoint_path: Path to a Lightning .ckpt file. If None, resolves
            via VLT's weight system.
        device: Device string ("cpu", "cuda", "auto"). "auto" picks CUDA if available.
    """

    def __init__(
        self,
        checkpoint_path: Path | str | None = None,
        device: str = "auto",
    ):
        if checkpoint_path is None:
            from sltk.extraction.weights import get_weight_path

            checkpoint_path = get_weight_path("segmentor")
            if checkpoint_path is None:
                raise FileNotFoundError(
                    "Segmentation checkpoint not found. Set SLTK_SEGMENTOR_CHECKPOINT "
                    "or place segmentor_v2.ckpt in sltk/weights/segmentor/"
                )

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

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

        # Load checkpoint to extract hyperparameters
        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        hp = ckpt["hyper_parameters"]
        model_config = hp["model_config"]

        # Extract inference parameters from checkpoint
        self.window_size = 25
        self.temp_conv_kernel = model_config["encoder"].get("temp_conv_kernel", 3)
        self._try_detect_window_size(ckpt)

        # Instantiate model and load state dict manually
        from sltk.segmentation.model import Transformer_Encoder

        dummy_dataset = self._make_dummy_dataset(model_config)
        self.model = Transformer_Encoder(
            model_config=model_config,
            train_batch_size=hp.get("train_batch_size", 64),
            val_batch_size=hp.get("val_batch_size", 64),
            dataset=dummy_dataset,
            model_dir="",
            fps=hp.get("fps", 50),
        )
        self.model.load_state_dict(ckpt["state_dict"])
        self.model.to(self.device)
        self.model.eval()

        logger.info(
            "Segmentation model loaded on %s (window_size=%d, temp_conv_kernel=%d)",
            self.device,
            self.window_size,
            self.temp_conv_kernel,
        )

    def _try_detect_window_size(self, ckpt: dict):
        """Try to detect window_size from checkpoint metadata."""
        hp = ckpt.get("hyper_parameters", {})
        dataset_cfg = hp.get("dataset", {})
        if isinstance(dataset_cfg, dict) and "window_size" in dataset_cfg:
            self.window_size = dataset_cfg["window_size"]
            return
        self.window_size = 25

    def _make_dummy_dataset(self, model_config: dict):
        """Create a minimal dummy dataset for model initialization."""

        class DummyTrain:
            input_size = 192
            gloss_vocab = list(range(11265))
            fps = 50
            subsample_pose = 1

            def get_class_weight(self, weight_type, converge_value):
                return torch.tensor([1.0, 1.0, 1.0])

        class DummyDataset:
            train = DummyTrain()

        return DummyDataset()

    @torch.no_grad()
    def predict(
        self,
        features: torch.Tensor,
        include_begin: bool = True,
    ) -> torch.Tensor:
        """Run segmentation on a single feature tensor.

        Args:
            features: Tensor of shape (num_frames, 192).
            include_begin: If True, keep BEGIN (2) labels. If False, merge to IN (1).

        Returns:
            Tensor of shape (num_frames,) with integer labels (0=OUT, 1=IN, 2=BEGIN).
        """
        features = features.to(self.device)
        temp_pad = (self.temp_conv_kernel - 1) // 2
        window = self.window_size
        num_frames = features.shape[0]

        # Split into windows
        a_overlap = num_frames % window
        overlap_feature = features[-a_overlap:] if a_overlap > 0 else None

        chunks = []
        for start in range(0, num_frames - (a_overlap if a_overlap > 0 else 0), window):
            end = start + window
            if end <= num_frames:
                chunks.append(features[start:end])

        if overlap_feature is not None:
            chunks.append(overlap_feature)

        if not chunks:
            chunks = [features]

        src_len = torch.tensor([c.shape[0] for c in chunks], device=self.device)

        # Pad for temporal convolution
        src_padded_list = []
        for src in chunks:
            left = src[0:1, :].expand(temp_pad, -1)
            right = src[-1:, :].expand(temp_pad, -1)
            src_padded = torch.cat([left, src, right], dim=0)
            src_padded_list.append(src_padded)

        src = pad_sequence(src_padded_list, batch_first=True, padding_value=0.0).to(torch.float32)
        src = src.to(self.device)

        max_src_len = src.shape[1] - 2 * temp_pad
        src_mask = (
            (torch.arange(max_src_len, device=self.device).unsqueeze(0) < src_len.unsqueeze(1))
            .bool()
            .unsqueeze(1)
        )

        batch = {
            "src": src,
            "src_length": src_len,
            "src_mask": src_mask,
        }

        sign_level, _ = self.model.model_forward(**batch)

        # Stitch: all full windows + the remainder crop
        if a_overlap > 0 and len(chunks) > 1:
            sign_crop = sign_level[-1, -a_overlap:]
            sign_level = sign_level[:-1].flatten(0, 1)
            sign_level = torch.cat((sign_level, sign_crop), dim=0)
        else:
            sign_level = sign_level.flatten(0, 1)[:num_frames]

        sign_level = sign_level.detach().cpu()
        _, labels = torch.max(sign_level, dim=-1)

        # Post-process
        labels = label_fixes(labels)
        if not include_begin:
            labels = merge_begin_tags(labels)

        return labels

    def run_h5(
        self,
        input_path: Path | str,
        include_begin: bool = True,
    ) -> dict[str, torch.Tensor]:
        """Run inference on H5 file(s).

        Args:
            input_path: Path to a single .h5 file or a directory of .h5 files.
            include_begin: Whether to include BEGIN labels in output.

        Returns:
            Dict mapping video name to prediction tensor.
        """
        input_path = Path(input_path)
        if input_path.is_file():
            features_dict = {input_path.stem: h5_to_features(input_path)}
        elif input_path.is_dir():
            features_dict = load_h5_directory(input_path)
        else:
            raise FileNotFoundError(f"Input path not found: {input_path}")

        predictions = {}
        for name, features in features_dict.items():
            logger.info("Segmenting %s (%d frames)", name, features.shape[0])
            labels = self.predict(features, include_begin=include_begin)
            predictions[name] = labels

        return predictions


# ---------------------------------------------------------------------------
# Thread-safe singleton
# ---------------------------------------------------------------------------

_runner: InferenceRunner | None = None
_runner_lock = threading.Lock()


def get_runner(
    checkpoint_path: Path | str | None = None,
    device: str = "auto",
) -> InferenceRunner:
    """Get or create the singleton InferenceRunner.

    Thread-safe with double-checked locking.
    """
    global _runner
    if _runner is not None:
        return _runner
    with _runner_lock:
        if _runner is not None:
            return _runner
        _runner = InferenceRunner(checkpoint_path=checkpoint_path, device=device)
        return _runner


# ---------------------------------------------------------------------------
# Convenience functions
# ---------------------------------------------------------------------------


def segment_h5(
    input_path: Path | str,
    output_path: Path | str | None = None,
    output_format: OutputFormat = OutputFormat.JSON,
    fps: float = 25.0,
    include_begin: bool = True,
    checkpoint_path: Path | str | None = None,
    media_path: Path | str | None = None,
) -> dict[str, torch.Tensor]:
    """One-shot H5 segmentation.

    Args:
        input_path: Path to H5 file or directory.
        output_path: Optional path to save results.
        output_format: Format for saved results.
        fps: Frames per second for timestamp conversion.
        include_begin: Whether to include BEGIN labels.
        checkpoint_path: Optional explicit checkpoint path.
        media_path: Optional media file path (used for ELAN output).

    Returns:
        Dict mapping video name to prediction tensor.
    """
    runner = get_runner(checkpoint_path=checkpoint_path)
    predictions = runner.run_h5(input_path, include_begin=include_begin)

    if output_path is not None:
        save_predictions(
            predictions,
            Path(output_path),
            output_format,
            fps=fps,
            media_path=media_path,
        )

    return predictions


def segment_features(
    features: torch.Tensor,
    include_begin: bool = True,
    checkpoint_path: Path | str | None = None,
) -> torch.Tensor:
    """Segment a pre-loaded feature tensor.

    Args:
        features: Tensor of shape (num_frames, 192).
        include_begin: Whether to include BEGIN labels.
        checkpoint_path: Optional explicit checkpoint path.

    Returns:
        Tensor of shape (num_frames,) with BIO labels.
    """
    runner = get_runner(checkpoint_path=checkpoint_path)
    return runner.predict(features, include_begin=include_begin)
