import json
import os
from pathlib import Path

import torch
from transformers import (
    AutoModelForSequenceClassification
)


class EndpointHandler:
    def __init__(self, path=""):
        folder = Path(path)

        self.device = (
            "cuda" if torch.cuda.is_available()
            else "cpu"
        )

        self.model = (
            AutoModelForSequenceClassification
            .from_pretrained(
                str(folder),
                trust_remote_code=True,
                local_files_only=True,
            )
        )

        self.model.to(self.device).eval()

        self.thresholds = {}

        for name in ("dense", "moe"):
            data = json.loads(
                (folder / name / "metrics.json")
                .read_text()
            )

            self.thresholds[name] = {
                key: float(value["threshold"])
                for key, value
                in data["thresholds"].items()
            }

        # Preload gated encoder and fail startup if
        # required credentials or model access are missing.
        self.model._get_encoder()

    @torch.inference_mode()
    def __call__(self, data):
        texts = data.get("inputs")
        params = data.get("parameters") or {}

        if isinstance(texts, str):
            texts = [texts]

        if not isinstance(texts, list):
            raise ValueError(
                "inputs must be a string or list of strings"
            )

        head = params.get("head", "moe")
        policy = params.get(
            "threshold_id", "T0_fixed_0_5"
        )

        if head not in ("dense", "moe"):
            raise ValueError("Unsupported head")

        if policy not in self.thresholds[head]:
            raise ValueError("Unsupported threshold")

        output = self.model(
            texts=texts,
            head=head,
        )

        scores = torch.softmax(
            output.logits, dim=-1
        )[:, 1].detach().cpu().tolist()

        threshold = self.thresholds[head][policy]

        return {
            "model": "selimk92/AstraeaGuard-V2",
            "head": head,
            "threshold_id": policy,
            "threshold": threshold,
            "status": "research_only",
            "blocking_enabled": False,
            "results": [
                {
                    "label": (
                        "ATTACK"
                        if score >= threshold
                        else "BENIGN"
                    ),
                    "attack_score": float(score),
                    "flagged": bool(
                        score >= threshold
                    ),
                }
                for score in scores
            ],
        }