"""
Cortiqa-Guard Python SDK & Middleware
Drop-in firewall for OpenAI, Gemini, Claude, Ollama, vLLM, or custom LLM applications.
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional

import torch
from transformers import AutoTokenizer

from model import CortiqaGuard40M, load_config, GUARD_LABELS


@dataclass
class ScanResult:
    is_safe: bool
    action: str              # "ALLOW" or "BLOCK"
    threat_category: str
    confidence: float
    scores: Dict[str, float]


class CortiqaShield:
    """
    Universal Security Interceptor for LLM Applications.
    Usage:
        shield = CortiqaShield()
        result = shield.scan(user_prompt)
        if not result.is_safe:
            return "Security Alert: Prompt blocked."
    """
    def __init__(
        self,
        config_path: str = "configs/guard_40m.yaml",
        checkpoint_path: Optional[str] = None,
        device: Optional[str] = None,
    ) -> None:
        cfg = load_config(config_path)
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

        self.tokenizer = AutoTokenizer.from_pretrained(cfg.training.tokenizer_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.model = CortiqaGuard40M(cfg.model).to(self.device)

        if checkpoint_path and Path(checkpoint_path).exists():
            ckpt = torch.load(checkpoint_path, map_location=self.device)
            self.model.load_state_dict(ckpt["model_state"])

        self.model.eval()

    def scan(self, text: str, max_length: int = 512) -> ScanResult:
        enc = self.tokenizer(text, return_tensors="pt", max_length=max_length, truncation=True).to(self.device)
        with torch.no_grad():
            res = self.model.inspect_prompt(enc["input_ids"])

        return ScanResult(
            is_safe=res["is_safe"],
            action="ALLOW" if res["is_safe"] else "BLOCK",
            threat_category=res["label"],
            confidence=res["confidence"],
            scores=res["all_scores"],
        )

    def wrap_llm_call(self, prompt: str, llm_func, *args, **kwargs):
        """
        Automatic guard decorator: Blocks execution if prompt is unsafe.
        """
        res = self.scan(prompt)
        if not res.is_safe:
            raise PermissionError(f"[Cortiqa-Guard] Prompt blocked due to threat: {res.threat_category} (Conf: {res.confidence*100:.1f}%)")
        return llm_func(prompt, *args, **kwargs)