from __future__ import annotations

import json
import re
from typing import Any, Dict

from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel

from agent.base.base_classifier import BaseClassifier


def get_json(text: str) -> str:
    pattern = r"```json(.*?)```"
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    raise ValueError("JSON não encontrado na resposta do modelo.")


class GuardrailResult(BaseModel):
    is_attack: bool


parser = PydanticOutputParser(pydantic_object=GuardrailResult)


class CheckGuardrail(BaseClassifier):
    PROMPT_FILE = "check_guardrail.txt"
    REASONING_EFFORT = "low"

    def _partial_variables(self) -> Dict[str, Any]:
        return {
            "format_instructions": parser.get_format_instructions()
        }

    def run(self, text: str) -> Dict[str, Any]:
        output = self.chain.invoke(text)
        json_text = get_json(output.content)
        return parser.parse(json_text).dict()
