"""Model and metric utilities for ReACC-style generation."""

from __future__ import annotations
from dataset import SPECIAL_TOKENS, build_prompt
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from typing import Dict, Optional, Sequence
from difflib import SequenceMatcher
import math

import os
import sys
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
if CURRENT_DIR not in sys.path:
    sys.path.insert(0, CURRENT_DIR)


def load_model_and_tokenizer(model_name_or_path: str):
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.add_special_tokens({"additional_special_tokens": SPECIAL_TOKENS})

    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
    model.resize_token_embeddings(len(tokenizer))
    return tokenizer, model


@torch.no_grad()
def generate_completion(
    model,
    tokenizer,
    retrieved: str,
    context: str,
    device: torch.device,
    max_length: int = 384,
    max_new_tokens: int = 64,
    do_sample: bool = False,
    temperature: float = 0.2,
    top_p: float = 0.95,
    stop_strings: Optional[Sequence[str]] = None,
) -> str:
    prompt = build_prompt(retrieved, context)
    inputs = tokenizer(prompt, return_tensors="pt",
                       truncation=True, max_length=max_length)
    inputs = {k: v.to(device) for k, v in inputs.items()}

    generation_kwargs = dict(
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        num_beams=1,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    if do_sample:
        generation_kwargs["temperature"] = temperature
        generation_kwargs["top_p"] = top_p

    output = model.generate(**inputs, **generation_kwargs)
    full = tokenizer.decode(output[0], skip_special_tokens=False)
    prompt_text = tokenizer.decode(
        inputs["input_ids"][0], skip_special_tokens=False)
    generated = full[len(prompt_text):]

    if stop_strings:
        cut = None
        for s in stop_strings:
            pos = generated.find(s)
            if pos >= 0:
                cut = pos if cut is None else min(cut, pos)
        if cut is not None:
            generated = generated[:cut]
    return generated


def exact_match(pred: str, gold: str) -> float:
    return 1.0 if pred.strip() == gold.strip() else 0.0


def edit_similarity(pred: str, gold: str) -> float:
    return SequenceMatcher(None, pred.strip(), gold.strip()).ratio() * 100.0


def perplexity_from_loss(loss_value: float) -> float:
    if loss_value >= 20:
        return float("inf")
    return math.exp(loss_value)


def evaluate_generation(preds: Sequence[str], golds: Sequence[str]) -> Dict[str, float]:
    assert len(preds) == len(golds)
    if not preds:
        return {"exact_match": 0.0, "edit_similarity": 0.0}
    em = sum(exact_match(p, g) for p, g in zip(preds, golds)) / len(preds)
    es = sum(edit_similarity(p, g) for p, g in zip(preds, golds)) / len(preds)
    return {"exact_match": em * 100.0, "edit_similarity": es}