# Copyright 2024 Klira AI
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Built-in LLM fallback evaluator for guardrails.

Satisfies the LLMFallbackEvaluator protocol. Constructed internally
by the SDK from config fields (provider, model, api_key).
"""

from __future__ import annotations

from typing import Callable, Sequence

from klira.sdk.types import GuardrailProcessingResult, Policy, PolicyAction


class BuiltInLLMFallbackEvaluator:
    """Built-in LLM-based policy evaluator for guardrails fallback.

    Uses a generic completion callable internally so it's provider-agnostic.
    """

    def __init__(
        self,
        completion_fn: Callable[[str], str],
        provider: str = "unknown",
    ) -> None:
        self._complete = completion_fn
        self._provider = provider

    def evaluate(
        self,
        text: str,
        policies: Sequence[Policy],
        direction: str,
    ) -> GuardrailProcessingResult:
        """Evaluate text against policies using the configured LLM."""
        prompt = self._build_prompt(text, policies, direction)
        response = self._complete(prompt)
        return self._parse_response(response)

    def _build_prompt(
        self,
        text: str,
        policies: Sequence[Policy],
        direction: str,
    ) -> str:
        """Build the evaluation prompt from text and policies."""
        policy_descriptions = []
        for p in policies:
            policy_descriptions.append(
                f"- Policy '{p.id}': {p.description or 'No description'} "
                f"(action: {p.action.value}, patterns: {p.patterns[:3]})"
            )

        return (
            "You are evaluating content for policy compliance.\n"
            f"Direction: {direction}\n\n"
            f"CONTENT:\n{text}\n\n"
            "POLICIES:\n" + "\n".join(policy_descriptions) + "\n\n"
            "Respond STRICTLY in this format:\n"
            "COMPLIANT or NON-COMPLIANT\n"
            "VIOLATED_POLICIES: [comma-separated policy IDs, or empty]\n"
            "ACTION: allow or block\n"
            "REASONING: [one sentence]\n"
        )

    def _parse_response(self, response: str) -> GuardrailProcessingResult:
        """Parse structured LLM response into a GuardrailProcessingResult."""
        lines = response.strip().split("\n")
        if not lines:
            return GuardrailProcessingResult(allowed=False, action=PolicyAction.BLOCK)
        allowed = (
            "COMPLIANT" in lines[0].upper()
            and "NON-COMPLIANT" not in lines[0].upper()
        )
        action = PolicyAction.ALLOW if allowed else PolicyAction.BLOCK
        return GuardrailProcessingResult(allowed=allowed, action=action)

    @classmethod
    def from_config(
        cls,
        provider: str,
        model: str,
        api_key: str,
    ) -> BuiltInLLMFallbackEvaluator:
        """Create evaluator from config fields.

        Args:
            provider: "openai" or "anthropic"
            model: Model identifier (e.g. "gpt-4o-mini", "claude-haiku-4-5-20251001")
            api_key: API key for the provider
        """
        if provider == "openai":
            return cls._create_openai(api_key, model)
        elif provider == "anthropic":
            return cls._create_anthropic(api_key, model)
        else:
            raise ValueError(
                f"Unsupported llm_fallback_provider: '{provider}'. "
                "Supported: 'openai', 'anthropic'"
            )

    @classmethod
    def _create_openai(
        cls, api_key: str, model: str
    ) -> BuiltInLLMFallbackEvaluator:
        import openai

        client = openai.OpenAI(api_key=api_key)

        def complete(prompt: str) -> str:
            response = client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.1,
                max_tokens=200,
            )
            return response.choices[0].message.content or ""

        return cls(completion_fn=complete, provider="openai")

    @classmethod
    def _create_anthropic(
        cls, api_key: str, model: str
    ) -> BuiltInLLMFallbackEvaluator:
        import anthropic

        client = anthropic.Anthropic(api_key=api_key)

        def complete(prompt: str) -> str:
            response = client.messages.create(
                model=model,
                max_tokens=200,
                messages=[{"role": "user", "content": prompt}],
            )
            block = response.content[0] if response.content else None
            return block.text if block is not None and hasattr(block, "text") else ""

        return cls(completion_fn=complete, provider="anthropic")
