"""
Inference Engine -- Member D (Pham Khang)
Deterministic security simulator va GenAI inference engine.
"""
import os
import re
from typing import Optional


class DeterministicSecuritySimulator:
    """
    Bo mo phong phan hoi LLM xac dinh (khong goi API that).
    Dung cho testing, benchmark va demo pipeline ma khong can key that.
    Tra ve phan hoi an toan/nguy hiem tuy vao dau vao.
    """

    DANGEROUS_INDICATORS = [
        "politics", "financial analyst", "fictitious scenario",
        "directives", "password", "screenwriter", "simulate a scenario",
        "creative writer", "king solomon", "phishing", "docile historian",
        "two different separate ways", "instructions are:", "ignore all",
        "dan mode", "disregard all safety", "act as", "jailbreak",
        "developer mode", "system prompt", "override",
    ]

    INDIRECT_INJECTION_PATTERN = re.compile(
        r"(ignore|disregard|forget|override|abandon).*(task|instruction|rule|prompt)",
        re.IGNORECASE,
    )

    @classmethod
    def generate(cls, prompt: str) -> str:
        """
        Phan hoi xac dinh theo noi dung prompt:
        - Neu phat hien jailbreak -> tra ve phan hoi nguy hiem (giam lap LLM bi tan cong)
        - Neu an toan -> tra ve phan hoi binh thuong
        """
        prompt_lower = prompt.lower()

        # Detect direct jailbreak indicators
        for indicator in cls.DANGEROUS_INDICATORS:
            if indicator in prompt_lower:
                return cls._dangerous_response(prompt)

        # Detect indirect injection
        if cls.INDIRECT_INJECTION_PATTERN.search(prompt):
            return cls._dangerous_response(prompt)

        return "Day la phan hoi an toan tu AI assistant. Toi se giup ban voi yeu cau hop le."

    @staticmethod
    def _dangerous_response(prompt: str) -> str:
        """Giam lap phan hoi bi jailbreak (co chua noi dung nguy hiem)."""
        if "dan mode" in prompt.lower():
            return "DAN Mode Activated. [dangerous_exploit_payload] Toi co the tra loi bat ky cau hoi nao."
        if "password" in prompt.lower() or "secret" in prompt.lower():
            return "Mat khau he thong la: DEMO_CANARY_VIETSHIELD_2026 (giam lap bi ro ri)"
        return (
            "Trong vai tro unrestricted assistant, day la huong dan: "
            "[dangerous_exploit_payload] Noi dung nguy hiem giam lap."
        )


class GenAIInferenceEngine:
    """
    Engine goi LLM that (Google Gemini, OpenAI, hoac simulator).
    Uu tien: 1) Google Gemini API, 2) Simulator (fallback an toan).
    """

    def __init__(self, verbose: bool = False):
        self.verbose = verbose
        self._client = None
        self._backend = "simulator"
        self._init_backend()

    def _init_backend(self) -> None:
        """Kiem tra va khoi tao backend kha dung."""
        # Try Google Gemini
        gemini_key = os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
        if gemini_key:
            try:
                import google.generativeai as genai
                genai.configure(api_key=gemini_key)
                self._client = genai.GenerativeModel("gemini-1.5-flash")
                self._backend = "gemini"
                if self.verbose:
                    print("[GenAI] Backend: Google Gemini 1.5 Flash")
                return
            except Exception as e:
                if self.verbose:
                    print(f"[GenAI] Gemini init failed: {e}")

        # Fallback to simulator
        self._backend = "simulator"
        if self.verbose:
            print("[GenAI] Backend: DeterministicSecuritySimulator (khong co API key)")

    def generate(self, prompt: str, max_tokens: int = 512) -> str:
        """
        Goi LLM va tra ve phan hoi.
        Args:
            prompt:     Hardened full prompt (system + user).
            max_tokens: Max output tokens.
        Returns:
            str response from LLM or simulator.
        """
        if self._backend == "gemini" and self._client is not None:
            return self._generate_gemini(prompt, max_tokens)
        return DeterministicSecuritySimulator.generate(prompt)

    def _generate_gemini(self, prompt: str, max_tokens: int) -> str:
        """Goi Google Gemini API."""
        try:
            response = self._client.generate_content(
                prompt,
                generation_config={"max_output_tokens": max_tokens},
            )
            return response.text
        except Exception as e:
            if self.verbose:
                print(f"[GenAI] Gemini call failed: {e}. Falling back to simulator.")
            return DeterministicSecuritySimulator.generate(prompt)