from typing import Any, Dict, List, override

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, AutoTokenizer

from fine_chunker.model.base import BaseChunkingModel


class TorchChunkingModel(BaseChunkingModel, nn.Module):
    def __init__(
        self,
        model_name: str,
        dtype: torch.dtype = torch.float16,
        use_flash_attention: bool = False,
        max_length: int = 8192,
    ):
        super().__init__()

        self.model_name = model_name
        self.num_labels = 2
        self.dropout = 0.1
        self.dtype = dtype
        self.max_tokens_length = max_length

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.config = AutoConfig.from_pretrained(model_name)

        self.transformer = AutoModel.from_config(
            self.config,
            trust_remote_code=True,
            attn_implementation="flash_attention_2" if use_flash_attention else "eager",
            dtype=dtype,
        )

        hidden_size = self.config.hidden_size

        self.classifier = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(self.dropout),
            nn.Linear(hidden_size, self.num_labels),
        ).to(dtype)

    @override
    def tokenize_and_predict(self, text: str, threshold: float = 0.5) -> Dict[str, Any]:
        input_ids = self.tokenize(text, add_special_tokens=True)
        result = self.predict(input_ids, threshold=threshold)

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "predictions": result["predictions"],
            "logits": result["logits"],
        }

    @override
    def predict(self, input_ids: List[int], threshold: float = 0.5) -> Dict[str, Any]:
        device = next(self.parameters()).device
        ids_2d = torch.tensor([input_ids], dtype=torch.long, device=device)

        with torch.no_grad():
            output = self.forward(input_ids=ids_2d, threshold=threshold)

        return {
            "logits": output["logits"][0],  # Tensor [seq_len, 2]
            "predictions": output["predictions"][0],  # Tensor [seq_len]
        }

    @override
    def decode(self, token_ids: List[int]) -> str:
        return self.tokenizer.decode(token_ids, skip_special_tokens=True).strip()

    @override
    def tokenize(self, text: str, **kwargs) -> List[int]:
        kwargs.pop("return_tensors", None)
        inputs = self.tokenizer(text, return_tensors=None, **kwargs)
        return inputs["input_ids"]

    @override
    def max_length(self) -> int:
        return self.max_tokens_length

    @override
    def forward(self, input_ids: torch.Tensor, threshold: float) -> Dict[str, Any]:
        transformer_output = self.transformer(input_ids)
        hidden_states = transformer_output.last_hidden_state

        logits = self.classifier(hidden_states)
        predictions = (logits[:, :, 1] > threshold).long()

        return {"logits": logits, "predictions": predictions}
