from typing import Any, Dict, List

import torch
from transformers import BertTokenizerFast

from model import BertForTokenAndSequenceJointClassification, TOKEN_TAGS, SEQUENCE_TAGS

MAX_LENGTH = 512
# Special/non-technique tags at the front of TOKEN_TAGS (see model.py).
IGNORED_TAG_IDS = {0, 1}  # "<PAD>", "O"
# This repo ships weights + config only, no tokenizer files (vocab.txt /
# tokenizer_config.json). config.json's vocab_size (28996) matches
# bert-base-cased exactly, so the tokenizer is loaded from the public base
# checkpoint instead of the local repo path.
TOKENIZER_BASE = "bert-base-cased"


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.tokenizer = BertTokenizerFast.from_pretrained(TOKENIZER_BASE)
        self.model = BertForTokenAndSequenceJointClassification.from_pretrained(path)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        text = data.get("inputs")
        if not isinstance(text, str) or not text:
            return {"error": "expected data.inputs to be a non-empty string"}

        encoding = self.tokenizer(
            text,
            return_offsets_mapping=True,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_LENGTH,
        )
        offset_mapping = encoding.pop("offset_mapping")[0].tolist()

        with torch.no_grad():
            outputs = self.model(**encoding)

        # token_logits: [1, seq_len, 20] at inference (labels=None), see model.py forward().
        token_probs = torch.softmax(outputs.token_logits[0], dim=-1)
        token_tag_ids = token_probs.argmax(dim=-1).tolist()
        token_tag_probs = token_probs.max(dim=-1).values.tolist()

        sequence_probs = torch.softmax(outputs.sequence_logits[0], dim=-1).tolist()
        sequence_tag_id = int(torch.argmax(outputs.sequence_logits[0]).item())

        # Merge consecutive TOKENS (by index, not character offset) sharing
        # the same non-ignored tag into one span. PTC/this model has no BIO
        # scheme: adjacent same-tag tokens belong to the same span by
        # construction (see GUIDA_MIGRAZIONE, §1). Character offsets are NOT
        # a valid adjacency test here: a word boundary space makes the next
        # token's start > the previous token's end even though they're the
        # same run of tokens, so merging must key off token index, not char
        # offset. Special tokens ([CLS]/[SEP]/[PAD]) have offset (0, 0) and
        # are skipped naturally since they never carry a real technique tag.
        spans: List[Dict[str, Any]] = []
        current = None
        for i, (tag_id, prob, (start, end)) in enumerate(zip(token_tag_ids, token_tag_probs, offset_mapping)):
            is_technique = tag_id not in IGNORED_TAG_IDS and end > start
            if (
                is_technique
                and current is not None
                and current["tag_id"] == tag_id
                and i == current["last_index"] + 1
            ):
                current["end"] = end
                current["probs"].append(prob)
                current["last_index"] = i
            else:
                if current is not None:
                    spans.append(current)
                current = (
                    {"tag_id": tag_id, "start": start, "end": end, "probs": [prob], "last_index": i}
                    if is_technique
                    else None
                )
        if current is not None:
            spans.append(current)

        out_spans = [
            {
                "technique": TOKEN_TAGS[s["tag_id"]],
                "start": s["start"],
                "end": s["end"],
                "text": text[s["start"] : s["end"]],
                "confidence": sum(s["probs"]) / len(s["probs"]),
            }
            for s in spans
        ]

        return {
            "spans": out_spans,
            "sequence_label": SEQUENCE_TAGS[sequence_tag_id],
            "sequence_probs": {SEQUENCE_TAGS[i]: p for i, p in enumerate(sequence_probs)},
            "truncated": len(encoding["input_ids"][0]) >= MAX_LENGTH,
        }