from __future__ import annotations

from typing import Any

import torch
from torch import nn
from transformers import AutoModel

from tiny_router.config import RouterModelConfig
from tiny_router.constants import HEAD_LABELS


class TinyRouterModel(nn.Module):
    def __init__(self, config: RouterModelConfig):
        super().__init__()
        self.config = config
        self.encoder = AutoModel.from_pretrained(config.encoder_name)
        hidden_size = self.encoder.config.hidden_size
        if config.pooling_type not in {"mean", "attention"}:
            raise ValueError(f"Unsupported pooling type: {config.pooling_type}")
        self.dropout = nn.Dropout(config.dropout)
        self.pooling_type = config.pooling_type
        if self.pooling_type == "attention":
            self.attention_pool = nn.Linear(hidden_size, 1)
        self.action_embedding = nn.Embedding(len(config.action_vocab), config.structured_hidden_dim)
        self.outcome_embedding = nn.Embedding(len(config.outcome_vocab), config.structured_hidden_dim)
        self.recency_projection = nn.Sequential(
            nn.Linear(1, config.recency_embed_dim),
            nn.ReLU(),
            nn.Linear(config.recency_embed_dim, config.recency_embed_dim),
        )
        structured_dim = (config.structured_hidden_dim * 2) + config.recency_embed_dim
        canonical_head_order = list(HEAD_LABELS.keys())
        extra_heads = [head for head in config.label_maps if head not in HEAD_LABELS]
        self.head_order = [head for head in canonical_head_order if head in config.label_maps] + extra_heads
        self.head_dependencies = {
            head: self.head_order[:idx]
            for idx, head in enumerate(self.head_order)
        }
        self.dependency_projections = nn.ModuleDict()
        self.heads = nn.ModuleDict(
            self._build_heads(hidden_size + structured_dim)
        )
        self.loss_fns = {head: nn.CrossEntropyLoss() for head in config.label_maps}

    def _build_heads(self, base_dim: int) -> dict[str, nn.Linear]:
        heads: dict[str, nn.Linear] = {}
        for head, labels in self.config.label_maps.items():
            input_dim = base_dim
            if self.config.use_head_dependencies and self.head_dependencies[head]:
                dependency_input_dim = sum(
                    len(self.config.label_maps[source])
                    for source in self.head_dependencies[head]
                )
                self.dependency_projections[head] = nn.Sequential(
                    nn.Linear(dependency_input_dim, self.config.dependency_hidden_dim),
                    nn.Tanh(),
                    nn.Dropout(self.config.dropout),
                )
                input_dim += self.config.dependency_hidden_dim
            heads[head] = nn.Linear(input_dim, len(labels))
        return heads

    def pool_hidden_state(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        if self.pooling_type == "attention":
            scores = self.attention_pool(hidden_state).squeeze(-1)
            scores = scores.masked_fill(attention_mask == 0, -1e4)
            weights = torch.softmax(scores, dim=-1).unsqueeze(-1)
            return torch.sum(hidden_state * weights, dim=1)
        mask = attention_mask.unsqueeze(-1).expand(hidden_state.size()).float()
        summed = torch.sum(hidden_state * mask, dim=1)
        counts = torch.clamp(mask.sum(dim=1), min=1e-6)
        return summed / counts

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        previous_action_id: torch.Tensor,
        previous_outcome_id: torch.Tensor,
        log_recency_seconds: torch.Tensor,
        has_interaction: torch.Tensor,
        has_recency: torch.Tensor | None = None,
        head_loss_weights: dict[str, float] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        pooled = self.pool_hidden_state(outputs.last_hidden_state, attention_mask)
        pooled = self.dropout(pooled)

        structured_mask = has_interaction.float().unsqueeze(-1)
        recency_mask_source = has_interaction if has_recency is None else has_recency
        recency_mask = recency_mask_source.float().unsqueeze(-1)
        action_emb = self.action_embedding(previous_action_id) * structured_mask
        outcome_emb = self.outcome_embedding(previous_outcome_id) * structured_mask
        recency_emb = self.recency_projection(log_recency_seconds.unsqueeze(-1)) * recency_mask
        combined = torch.cat([pooled, action_emb, outcome_emb, recency_emb], dim=-1)

        logits = {}
        prior_head_probs: dict[str, torch.Tensor] = {}
        for head in self.head_order:
            head_inputs = [combined]
            if self.config.use_head_dependencies and self.head_dependencies[head]:
                dependency_input = torch.cat(
                    [prior_head_probs[source] for source in self.head_dependencies[head]],
                    dim=-1,
                )
                dependency_features = self.dependency_projections[head](dependency_input)
                head_inputs.append(dependency_features)
            classifier_input = torch.cat(head_inputs, dim=-1)
            head_logits = self.heads[head](classifier_input)
            logits[head] = head_logits
            # Keep the supervision for earlier heads local to their own loss while still
            # exposing their predictions to later heads.
            prior_head_probs[head] = torch.softmax(head_logits.detach(), dim=-1)
        result: dict[str, Any] = {"logits": logits}

        total_loss = None
        for head in self.config.label_maps:
            label_key = f"labels_{head}"
            if label_key not in kwargs:
                continue
            loss = self.loss_fns[head](logits[head], kwargs[label_key])
            if head_loss_weights:
                loss = loss * float(head_loss_weights.get(head, 1.0))
            total_loss = loss if total_loss is None else total_loss + loss
        result["loss"] = total_loss
        return result