# 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.

"""GuardrailsEngine — main entry point for guardrails processing.

Singleton with DI support for testing. Uses constructor injection
directly — no service locator.

Processing flow:
1. FastRulesEngine evaluates text against policies
2. DecisionRouter routes matches to decisions
3. Compliance audit spans created asynchronously
"""

from __future__ import annotations

import hashlib
import logging
import threading
from collections import OrderedDict
from typing import Sequence

from opentelemetry import trace

from klira.sdk.guardrails.decision import DecisionRouter
from klira.sdk.guardrails.fast_rules import FastRulesEngine
from klira.sdk.types import (
    GuardrailProcessingResult,
    LLMFallbackEvaluator,
    Policy,
    PolicyAction,
)

logger = logging.getLogger(__name__)


class _LLMFallbackCache:
    """Exact-match LRU cache for LLM fallback results."""

    def __init__(self, capacity: int = 1000) -> None:
        self._cache: OrderedDict[str, GuardrailProcessingResult] = OrderedDict()
        self._capacity = capacity
        self._lock = threading.Lock()

    def _key(self, text: str, policy_ids: frozenset[str], direction: str) -> str:
        return hashlib.md5(
            f"{text}:{sorted(policy_ids)}:{direction}".encode()
        ).hexdigest()

    def get(
        self, text: str, policy_ids: frozenset[str], direction: str
    ) -> GuardrailProcessingResult | None:
        key = self._key(text, policy_ids, direction)
        with self._lock:
            if key in self._cache:
                self._cache.move_to_end(key)
                return self._cache[key]
            return None

    def set(
        self,
        text: str,
        policy_ids: frozenset[str],
        direction: str,
        result: GuardrailProcessingResult,
    ) -> None:
        key = self._key(text, policy_ids, direction)
        with self._lock:
            if key in self._cache:
                self._cache.move_to_end(key)
            elif len(self._cache) >= self._capacity:
                self._cache.popitem(last=False)
            self._cache[key] = result


class GuardrailsEngine:
    """Main guardrails engine.

    Singleton for backward compatibility. Supports DI via constructor
    injection for testing — pass custom FastRulesEngine/DecisionRouter.
    """

    _instance: GuardrailsEngine | None = None
    _lock = threading.Lock()

    def __init__(
        self,
        policies: Sequence[Policy] | None = None,
        fast_rules: FastRulesEngine | None = None,
        decision_router: DecisionRouter | None = None,
        llm_fallback_enabled: bool = False,
        llm_evaluator: LLMFallbackEvaluator | None = None,
        llm_fallback_on_error: str = "allow",
    ) -> None:
        self._policies = list(policies) if policies else []
        self._fast_rules = fast_rules or FastRulesEngine()
        self._decision_router = decision_router or DecisionRouter()
        self._llm_fallback_enabled = llm_fallback_enabled
        self._llm_evaluator = llm_evaluator
        self._llm_fallback_on_error = llm_fallback_on_error
        self._llm_cache = _LLMFallbackCache() if llm_fallback_enabled else None
        self._tracer = trace.get_tracer("klira.guardrails")

        if llm_fallback_enabled and llm_evaluator is None:
            logger.warning(
                "llm_fallback_enabled=True but no llm_evaluator provided. "
                "LLM fallback will be skipped until an evaluator is set."
            )

    @classmethod
    def get_instance(cls) -> GuardrailsEngine:
        """Get the singleton instance. Creates one with defaults if needed."""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = cls()
        return cls._instance

    @classmethod
    def set_instance(cls, engine: GuardrailsEngine) -> None:
        """Set the singleton instance (used during Klira.init())."""
        with cls._lock:
            cls._instance = engine

    @classmethod
    def reset(cls) -> None:
        """Reset the singleton (used in testing and shutdown)."""
        with cls._lock:
            cls._instance = None

    @property
    def policies(self) -> list[Policy]:
        """Currently loaded policies."""
        return self._policies

    def set_policies(self, policies: Sequence[Policy]) -> None:
        """Update loaded policies (used after remote policy fetch)."""
        with self._lock:
            self._policies = list(policies)

    def check_input(
        self,
        text: str,
        domain: str | None = None,
        policy_set: str | None = None,
    ) -> GuardrailProcessingResult:
        """Check input text against guardrails policies.

        Args:
            text: The input text to check.
            domain: Optional domain for domain-based matching.
            policy_set: Optional policy set name to filter policies.

        Returns:
            GuardrailProcessingResult with the decision.
        """
        policies = self._filter_policies(policy_set)
        if not policies:
            logger.warning(
                "Guardrails check_input returning allow-by-default: "
                "no policies matched (policy_set=%s, total_policies=%d)",
                policy_set,
                len(self._policies),
            )
            with self._tracer.start_as_current_span(
                "klira.guardrails.input",
                attributes={
                    "klira.entity_type": "guardrails",
                    "klira.entity_name": "input",
                    "klira.compliance.direction": "inbound",
                    "klira.guardrails.policy_count": 0,
                    "klira.guardrails.empty_policy_set": True,
                    "klira.guardrails.decision": "allow",
                    "klira.guardrails.allowed": True,
                },
            ):
                return GuardrailProcessingResult(
                    allowed=True,
                    action=PolicyAction.ALLOW,
                )

        with self._tracer.start_as_current_span(
            "klira.guardrails.input",
            attributes={
                "klira.entity_type": "guardrails",
                "klira.entity_name": "input",
                "klira.compliance.direction": "inbound",
                "klira.guardrails.policy_count": len(policies),
            },
        ) as span:
            matches = self._fast_rules.evaluate(text, policies, domain)
            matched_policy_ids = {m.policy.id for m in matches}
            unmatched = (
                [p for p in policies if p.id not in matched_policy_ids]
                if self._llm_fallback_enabled
                else None
            )
            result = self._decision_router.route(
                matches,
                direction="inbound",
                text=text,
                unmatched_policies=unmatched,
                llm_evaluator=self._llm_evaluator if self._llm_fallback_enabled else None,
                llm_cache=self._llm_cache,
                llm_fallback_on_error=self._llm_fallback_on_error,
            )
            span.set_attribute("klira.guardrails.decision", result.action.value)
            span.set_attribute("klira.guardrails.allowed", result.allowed)
            return result

    def check_output(
        self,
        text: str,
        domain: str | None = None,
        policy_set: str | None = None,
    ) -> GuardrailProcessingResult:
        """Check output text against guardrails policies.

        Args:
            text: The output text to check.
            domain: Optional domain for domain-based matching.
            policy_set: Optional policy set name to filter policies.

        Returns:
            GuardrailProcessingResult with the decision.
        """
        policies = self._filter_policies(policy_set)
        if not policies:
            logger.warning(
                "Guardrails check_output returning allow-by-default: "
                "no policies matched (policy_set=%s, total_policies=%d)",
                policy_set,
                len(self._policies),
            )
            with self._tracer.start_as_current_span(
                "klira.guardrails.output",
                attributes={
                    "klira.entity_type": "guardrails",
                    "klira.entity_name": "output",
                    "klira.compliance.direction": "outbound",
                    "klira.guardrails.policy_count": 0,
                    "klira.guardrails.empty_policy_set": True,
                    "klira.guardrails.decision": "allow",
                    "klira.guardrails.allowed": True,
                },
            ):
                return GuardrailProcessingResult(
                    allowed=True,
                    action=PolicyAction.ALLOW,
                )

        with self._tracer.start_as_current_span(
            "klira.guardrails.output",
            attributes={
                "klira.entity_type": "guardrails",
                "klira.entity_name": "output",
                "klira.compliance.direction": "outbound",
                "klira.guardrails.policy_count": len(policies),
            },
        ) as span:
            matches = self._fast_rules.evaluate(text, policies, domain)
            matched_policy_ids = {m.policy.id for m in matches}
            unmatched = (
                [p for p in policies if p.id not in matched_policy_ids]
                if self._llm_fallback_enabled
                else None
            )
            result = self._decision_router.route(
                matches,
                direction="outbound",
                text=text,
                unmatched_policies=unmatched,
                llm_evaluator=self._llm_evaluator if self._llm_fallback_enabled else None,
                llm_cache=self._llm_cache,
                llm_fallback_on_error=self._llm_fallback_on_error,
            )
            span.set_attribute("klira.guardrails.decision", result.action.value)
            span.set_attribute("klira.guardrails.allowed", result.allowed)
            return result

    def _filter_policies(self, policy_set: str | None) -> list[Policy]:
        """Filter policies by policy_set if specified."""
        if policy_set is None:
            return self._policies
        return [p for p in self._policies if p.policy_set == policy_set]
