import json
from pathlib import Path

import numpy as np
from peft import PeftModel
from safetensors.numpy import load_file
from sentence_transformers import SentenceTransformer


class MuncheAIDetector:
    def __init__(self, model_dir: str | Path = "."):
        model_dir = Path(model_dir)
        self.encoder = SentenceTransformer("google/embeddinggemma-300m")
        self.encoder[0].auto_model = PeftModel.from_pretrained(
            self.encoder[0].auto_model,
            model_dir,
        )
        self.encoder.max_seq_length = 2048
        head = load_file(model_dir / "linear_head.safetensors")
        self.weight = head["linear.weight"].reshape(-1)
        self.bias = float(head["linear.bias"].item())
        calibration = json.loads(
            (model_dir / "calibration.json").read_text(encoding="utf-8")
        )
        self.human_max = calibration["thresholds"]["human_max"]
        self.ai_min = calibration["thresholds"]["ai_min"]

    def predict(self, text: str) -> dict:
        tokens = self.encoder.tokenizer.encode(text, add_special_tokens=False)
        if not 384 <= len(tokens) <= 2048:
            raise ValueError("Input must contain 384 to 2,048 EmbeddingGemma tokens.")
        embedding = self.encoder.encode(
            [text],
            normalize_embeddings=True,
            convert_to_numpy=True,
        )[0]
        logit = float(embedding @ self.weight + self.bias)
        score = (
            float(1.0 / (1.0 + np.exp(-logit)))
            if logit >= 0
            else float(np.exp(logit) / (1.0 + np.exp(logit)))
        )
        label = (
            "human"
            if score <= self.human_max
            else "llm"
            if score >= self.ai_min
            else "uncertain"
        )
        return {"label": label, "score": score, "tokens": len(tokens)}