import requests


class LLMJsonGuard:
    def __init__(self, api_key: str):
        if not api_key:
            raise ValueError("RapidAPI key is required")

        self.api_key = api_key
        self.base_url = "https://llm-json-sanitizer-schema-guard.p.rapidapi.com"
        self.headers = {
            "Content-Type": "application/json",
            "x-rapidapi-key": self.api_key,
            "x-rapidapi-host": "llm-json-sanitizer-schema-guard.p.rapidapi.com"
        }

    def sanitize(self, raw_output: str):
        if not raw_output:
            raise ValueError("raw_output is required")

        return self._request(
            "/api/json-repair/sanitize",
            {"raw_output": raw_output}
        )

    def guard(self, raw_output: str, schema: dict):
        if not raw_output or not schema:
            raise ValueError("raw_output and schema are required")

        return self._request(
            "/api/json-repair/guard",
            {
                "raw_output": raw_output,
                "schema": schema
            }
        )

    def _request(self, path: str, body: dict):
        response = requests.post(
            self.base_url + path,
            headers=self.headers,
            json=body,
            timeout=15
        )

        data = response.json()

        if not response.ok:
            message = (
                data.get("data", {})
                    .get("response", {})
                    .get("stage")
                or data.get("message")
                or "Request failed"
            )
            raise Exception(message)

        return data["data"]["response"]