"""Inference and evaluation for the encoder model."""

import argparse
import json
import torch
from pathlib import Path
from transformers import AutoTokenizer
from tqdm import tqdm
import sys

sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from src.encoder_model import EmojinizeEncoderModel
from src.encoder_dataset import load_encoder_datasets
from src import emoji_vocab
from src.validators import validate_annotation_json


def _load_model(checkpoint_path: str, device):
    """Load model using saved model_config.json so architecture is reconstructed exactly."""
    ckpt = Path(checkpoint_path)
    config_path = ckpt / "model_config.json"
    if config_path.exists():
        with open(config_path) as f:
            cfg = json.load(f)
    else:
        # Fallback for checkpoints saved before model_config.json was introduced
        cfg = {}
    model = EmojinizeEncoderModel(**cfg)
    model.load_state_dict(torch.load(ckpt / "model.pt", map_location=device, weights_only=True))
    return model


def _extract_word_level_spans(bio_labels, word_ids, offset_mapping, sentence):
    """Extract spans at word granularity, fixing both subword-splitting and trailing spaces.

    Strategy:
    - Group tokens by word_id (word_ids() from the tokenizer).
    - For each word, use only the FIRST subword's BIO prediction.
      Non-first subwords were trained with ignore_index=-100, so their
      predictions are meaningless and must be skipped.
    - Each word's char range = first_subword.char_start … last_subword.char_end,
      so spans are clean word-aligned boundaries with no trailing spaces.

    Returns:
        list of (tok_start, tok_end, char_start, char_end, span_text)
        where tok_start/end are token indices for span pooling.
    """
    from collections import defaultdict
    word_tokens = defaultdict(list)
    for tok_idx, wid in enumerate(word_ids):
        if wid is not None:
            word_tokens[wid].append(tok_idx)

    # Build word-level entries: (bio_label, first_tok, one_past_last_tok, char_s, char_e)
    words = []
    for wid in sorted(word_tokens.keys()):
        toks = word_tokens[wid]
        first_tok, last_tok = toks[0], toks[-1]
        label = bio_labels[first_tok]
        char_s = offset_mapping[first_tok][0]
        char_e = offset_mapping[last_tok][1]
        if char_s == -1 or char_e == -1:
            continue
        words.append((label, first_tok, last_tok + 1, char_s, char_e))

    # BIO span extraction on word-level labels; fix illegal I-without-B → B
    spans = []
    in_span = False
    s_tok = s_char = e_tok = e_char = None

    def _close():
        if s_tok is not None and s_char < e_char:
            span_text = sentence[s_char:e_char]
            # Strip trailing sentence punctuation (commas, periods, etc.) that tokenizer attaches
            # This ensures spans like "design," become "design" in output
            end_punct = ".,!?;:"
            trimmed_end = len(span_text)
            while trimmed_end > 0 and span_text[trimmed_end - 1] in end_punct:
                trimmed_end -= 1

            if trimmed_end > 0:
                # Adjust char boundaries to exclude trailing punctuation
                adjusted_e_char = s_char + trimmed_end
                spans.append((s_tok, e_tok, s_char, adjusted_e_char, span_text[:trimmed_end]))

    for label, tok_s, tok_e, char_s, char_e in words:
        if label == 1:  # B
            _close()
            s_tok, s_char = tok_s, char_s
            e_tok, e_char = tok_e, char_e
            in_span = True
        elif label == 2:  # I
            if not in_span:  # promote illegal I→B
                _close()
                s_tok, s_char = tok_s, char_s
                in_span = True
            e_tok, e_char = tok_e, char_e
        else:  # O
            _close()
            s_tok = None
            in_span = False

    _close()
    return spans


@torch.no_grad()
def infer_sentence(model, tokenizer, sentence, device):
    """Run the full pipeline on one sentence.

    Returns: {"marked": str, "annotations": list[{"span": str, "emojis": str}]}
    """
    enc = tokenizer(
        sentence,
        return_offsets_mapping=True,
        truncation=True,
        max_length=512,
    )
    input_ids = torch.tensor([enc["input_ids"]], dtype=torch.long, device=device)
    attn_mask = torch.tensor([enc["attention_mask"]], dtype=torch.long, device=device)
    offsets = enc["offset_mapping"]
    word_ids = enc.word_ids()  # list[int|None], one per token

    # ── Encoder + BIO ────────────────────────────────────────────────────────
    encoder_out = model.encoder(input_ids, attn_mask)
    hidden = encoder_out.last_hidden_state              # (1, seq_len, hidden)
    bio_logits = model.bio_head(hidden)                 # (1, seq_len, 3)
    # Viterbi decoding enforces valid B-I-O transitions
    decoded_tags = model.bio_head.decode(bio_logits, attn_mask)  # list[list[int]]
    seq_len = bio_logits.shape[1]
    # Pad back to seq_len with O (0) in case any padding positions were excluded
    raw = decoded_tags[0]
    bio_labels = raw + [0] * (seq_len - len(raw))

    # Word-level span extraction: fixes subword splitting and trailing spaces
    spans = _extract_word_level_spans(bio_labels, word_ids, offsets, sentence)
    if not spans:
        return {"marked": sentence, "annotations": []}

    # ── Batch all spans in ONE decoder call ──────────────────────────────────
    span_info_local = [(tok_s, tok_e) for tok_s, tok_e, *_ in spans]
    span_embs, enc_hidden, enc_mask = model._pool_spans(
        hidden, attn_mask, [span_info_local]
    )

    _, decoded_seqs = model.emoji_decoder(span_embs, enc_hidden, enc_mask)

    # ── Reconstruct output ───────────────────────────────────────────────────
    annotations = []
    for (_, _, cs, ce, span_text), token_ids in zip(spans, decoded_seqs):
        emoji_str = emoji_vocab.decode(token_ids)
        annotations.append({"span": span_text, "emojis": emoji_str})

    # Insert <span> tags in reverse order to preserve char indices
    marked = sentence
    for _, _, cs, ce, span_text in reversed(spans):
        marked = marked[:cs] + f"<span>{span_text}</span>" + marked[ce:]

    return {"marked": marked, "annotations": annotations}


def evaluate_dataset(model, tokenizer, device, dataset, num_samples=None):
    model.eval()
    n = min(num_samples or len(dataset), len(dataset))
    valid, invalid = 0, 0
    samples = []

    for idx in tqdm(range(n)):
        item = dataset.hf_dataset[idx]
        sentence = next(
            (m["content"] for m in item["prompt"] if m["role"] == "user"), ""
        )
        try:
            out = infer_sentence(model, tokenizer, sentence, device)
            ok, msg = validate_annotation_json(out, sentence)
            if ok:
                valid += 1
            else:
                invalid += 1
            if len(samples) < 5:
                samples.append({"sentence": sentence, "output": out, "valid": ok, "msg": msg})
        except Exception as e:
            invalid += 1
            if len(samples) < 5:
                samples.append({"sentence": sentence, "error": str(e), "valid": False})

    return {"valid": valid, "invalid": invalid, "samples": samples}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--data_path", default="data")
    parser.add_argument("--output_dir", default="eval_results")
    parser.add_argument("--num_samples", type=int, default=200)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    print("Loading model …")
    model = _load_model(args.checkpoint, device)
    model.to(device).eval()

    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint)

    _, eval_ds = load_encoder_datasets(args.data_path, tokenizer, max_length=512)

    print("Running evaluation …")
    results = evaluate_dataset(model, tokenizer, device, eval_ds, args.num_samples)

    total = results["valid"] + results["invalid"]
    print(f"\nValid  : {results['valid']}/{total}  ({results['valid']/total*100:.1f}%)")

    out_path = Path(args.output_dir) / "results.json"
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print(f"Saved → {out_path}")


if __name__ == "__main__":
    main()