from __future__ import annotations

import json
from contextlib import nullcontext
from pathlib import Path

import torch
from safetensors.torch import load_file, save_file
from torch import nn
from transformers import AutoModel, AutoTokenizer

from .config import ModelConfig


def resolve_torch_dtype(name: str, device: torch.device) -> torch.dtype:
    if name == "float32":
        return torch.float32
    if name == "float16":
        return torch.float16
    if name == "bfloat16":
        return torch.bfloat16
    if device.type != "cuda":
        return torch.float32
    if torch.cuda.is_bf16_supported():
        return torch.bfloat16
    return torch.float16


class RouterModel(nn.Module):
    """Frozen language-model encoder plus a small worker-selection head."""

    def __init__(
        self,
        config: ModelConfig,
        worker_ids: list[str],
        device: torch.device,
        backbone_source: str | Path | None = None,
    ):
        super().__init__()
        self.router_config = config
        self.worker_ids = list(worker_ids)
        self.device_ref = device
        dtype = resolve_torch_dtype(config.dtype, device)
        source = str(backbone_source or config.base_model)
        self.tokenizer = AutoTokenizer.from_pretrained(source, use_fast=True)
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "right"
        self.backbone = AutoModel.from_pretrained(source, torch_dtype=dtype)
        hidden_size = getattr(self.backbone.config, "hidden_size", None)
        if hidden_size is None:
            hidden_size = getattr(self.backbone.config, "d_model", None)
        if hidden_size is None:
            raise ValueError("Could not infer backbone hidden size")
        self.norm = nn.LayerNorm(hidden_size, dtype=torch.float32)
        self.dropout = nn.Dropout(config.dropout)
        self.classifier = nn.Linear(hidden_size, len(worker_ids), dtype=torch.float32)
        self.backbone.requires_grad_(not config.freeze_backbone)
        self.to(device)
        if config.freeze_backbone:
            self.backbone.eval()

    def train(self, mode: bool = True):
        super().train(mode)
        if self.router_config.freeze_backbone:
            self.backbone.eval()
        return self

    def _pool(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        if self.router_config.pooling == "mean":
            mask = attention_mask.unsqueeze(-1).to(hidden.dtype)
            return (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1)
        positions = torch.arange(hidden.shape[1], device=hidden.device).unsqueeze(0)
        last_indices = positions.masked_fill(attention_mask == 0, -1).max(dim=1).values
        batch_indices = torch.arange(hidden.shape[0], device=hidden.device)
        return hidden[batch_indices, last_indices]

    def encode_features(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        context = torch.no_grad() if self.router_config.freeze_backbone else nullcontext()
        with context:
            output = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
            pooled = self._pool(output.last_hidden_state, attention_mask)
        return self.norm(pooled.float())

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        features = self.encode_features(input_ids, attention_mask)
        return self.classifier(self.dropout(features))

    def head_state_dict(self) -> dict[str, torch.Tensor]:
        state = {}
        for prefix, module in (("norm", self.norm), ("classifier", self.classifier)):
            for name, value in module.state_dict().items():
                state[f"{prefix}.{name}"] = value.detach().cpu().contiguous()
        return state

    def load_head_state_dict(self, state: dict[str, torch.Tensor]) -> None:
        self.norm.load_state_dict(
            {key.removeprefix("norm."): value for key, value in state.items() if key.startswith("norm.")}
        )
        self.classifier.load_state_dict(
            {
                key.removeprefix("classifier."): value
                for key, value in state.items()
                if key.startswith("classifier.")
            }
        )

    def save_checkpoint(self, path: str | Path, metadata: dict | None = None) -> None:
        destination = Path(path)
        destination.mkdir(parents=True, exist_ok=True)
        payload = {
            "format_version": 1,
            "model": self.router_config.model_dump(),
            "worker_ids": self.worker_ids,
            "metadata": metadata or {},
        }
        (destination / "router_config.json").write_text(
            json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
        )
        save_file(self.head_state_dict(), destination / "router_head.safetensors")
        self.tokenizer.save_pretrained(destination / "tokenizer")
        if not self.router_config.freeze_backbone:
            self.backbone.save_pretrained(destination / "backbone", safe_serialization=True)

    @classmethod
    def from_checkpoint(
        cls, path: str | Path, device: torch.device | None = None
    ) -> RouterModel:
        source = Path(path)
        payload = json.loads((source / "router_config.json").read_text(encoding="utf-8"))
        config = ModelConfig.model_validate(payload["model"])
        target_device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        backbone_path = source / "backbone"
        model = cls(
            config,
            payload["worker_ids"],
            target_device,
            backbone_source=backbone_path if backbone_path.exists() else None,
        )
        tokenizer_path = source / "tokenizer"
        if tokenizer_path.exists():
            model.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
            if model.tokenizer.pad_token_id is None:
                model.tokenizer.pad_token = model.tokenizer.eos_token
            model.tokenizer.padding_side = "right"
        state = load_file(source / "router_head.safetensors", device="cpu")
        model.load_head_state_dict(state)
        return model

    def trainable_parameter_counts(self) -> tuple[int, int]:
        trainable = sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
        total = sum(parameter.numel() for parameter in self.parameters())
        return trainable, total