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

"""Decision router — routes guardrail matches to block/allow/augment decisions.

Creates compliance audit spans synchronously (microsecond overhead, reliable export).
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Sequence

from opentelemetry import trace

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

if TYPE_CHECKING:
    from klira.sdk.guardrails.engine import _LLMFallbackCache
    from klira.sdk.types import LLMFallbackEvaluator

logger = logging.getLogger(__name__)

_DEFAULT_VIOLATION_RESPONSE = "Your request has been blocked by a guardrails policy."


class DecisionRouter:
    """Routes guardrail evaluation results to decisions.

    Simple routing logic:
    1. If any blocking policies matched → block
    2. If augmentation policies matched → allow with guidelines
    3. If nothing matched → allow
    """

    def __init__(self) -> None:
        self._tracer = trace.get_tracer("klira.guardrails")

    def route(
        self,
        matches: list[PolicyMatch],
        direction: str = "inbound",
        text: str = "",
        unmatched_policies: list[Policy] | None = None,
        llm_evaluator: LLMFallbackEvaluator | None = None,
        llm_cache: _LLMFallbackCache | None = None,
        llm_fallback_on_error: str = "allow",
    ) -> GuardrailProcessingResult:
        """Route matches to a guardrails decision.

        Args:
            matches: List of policy matches from FastRulesEngine.
            direction: "inbound" or "outbound".
            text: Original text (needed for LLM fallback cache key).
            unmatched_policies: Policies that fast rules didn't match.
            llm_evaluator: Optional LLM fallback evaluator.
            llm_cache: Optional LLM fallback result cache.
            llm_fallback_on_error: "allow" or "block" on evaluator error.

        Returns:
            GuardrailProcessingResult with the decision.
        """
        blocking = [m for m in matches if m.policy.action == PolicyAction.BLOCK]
        augmenting = [m for m in matches if m.policy.action == PolicyAction.ALLOW]

        if blocking:
            result = self._make_block_decision(blocking)
            self._create_compliance_span(
                "klira.compliance.blocked", result, direction, matches
            )
            return result

        if augmenting and direction == "inbound":
            result = self._make_augment_decision(augmenting)
            self._create_compliance_span(
                "klira.compliance.augmented", result, direction, matches,
                layer="prompt_augmented",
            )
            return result

        # LLM fallback — when fast rules found nothing, ask LLM
        if llm_evaluator and unmatched_policies:
            policy_ids = frozenset(p.id for p in unmatched_policies)

            # Check cache first
            if llm_cache:
                cached = llm_cache.get(text, policy_ids, direction)
                if cached is not None:
                    self._create_compliance_span(
                        "klira.compliance.llm_fallback",
                        cached,
                        direction,
                        matches,
                        layer="llm_fallback",
                    )
                    return cached

            # Call LLM evaluator
            try:
                result = llm_evaluator.evaluate(
                    text, unmatched_policies, direction
                )
                if llm_cache:
                    llm_cache.set(text, policy_ids, direction, result)
                self._create_compliance_span(
                    "klira.compliance.llm_fallback",
                    result,
                    direction,
                    matches,
                    layer="llm_fallback",
                )
                return result
            except Exception:
                logger.warning(
                    "LLM fallback evaluation failed", exc_info=True
                )
                if llm_fallback_on_error == "block":
                    result = GuardrailProcessingResult(
                        allowed=False, action=PolicyAction.BLOCK
                    )
                    self._create_compliance_span(
                        "klira.compliance.llm_fallback",
                        result,
                        direction,
                        matches,
                        layer="llm_fallback",
                    )
                    return result
                # else: fall through to allow

        # No policies matched — allow through
        result = GuardrailProcessingResult(
            allowed=True,
            action=PolicyAction.ALLOW,
        )
        self._create_compliance_span(
            "klira.compliance.allowed", result, direction, matches
        )
        return result

    def _make_block_decision(
        self, blocking: list[PolicyMatch]
    ) -> GuardrailProcessingResult:
        """Create a block decision from blocking policy matches."""
        # Use the first blocking policy's violation response
        violation_response = _DEFAULT_VIOLATION_RESPONSE
        for match in blocking:
            if match.policy.violation_response:
                violation_response = match.policy.violation_response
                break

        # Get regulation from first blocking policy that has one
        regulation = None
        for match in blocking:
            if match.policy.regulation:
                regulation = match.policy.regulation
                break

        return GuardrailProcessingResult(
            allowed=False,
            action=PolicyAction.BLOCK,
            violation_response=violation_response,
            matched_policies=blocking,
            regulation=regulation,
        )

    def _make_augment_decision(
        self, augmenting: list[PolicyMatch]
    ) -> GuardrailProcessingResult:
        """Create an augment decision from allow policy matches."""
        # Collect all guidelines from matching policies
        guidelines: list[str] = []
        for match in augmenting:
            guidelines.extend(match.policy.guidelines)

        # Get regulation from first augmenting policy that has one
        regulation = None
        for match in augmenting:
            if match.policy.regulation:
                regulation = match.policy.regulation
                break

        return GuardrailProcessingResult(
            allowed=True,
            action=PolicyAction.ALLOW,
            guidelines=guidelines,
            matched_policies=augmenting,
            augmentation_applied=bool(guidelines),
            regulation=regulation,
        )

    def _create_compliance_span(
        self,
        span_name: str,
        result: GuardrailProcessingResult,
        direction: str,
        matches: Sequence[PolicyMatch],
        layer: str = "fast_rules",
    ) -> None:
        """Create a compliance audit span synchronously.

        Compliance spans are zero-duration marker spans with attributes —
        microsecond overhead, always reliably exported.
        """

        def _create_span() -> None:
            try:
                attrs: dict[str, str | bool | int] = {
                    "klira.entity_type": "compliance",
                    "klira.entity_name": span_name.split(".")[-1],
                    "klira.compliance.direction": direction,
                    "klira.compliance.decision.allowed": result.allowed,
                    "klira.compliance.decision.action": result.action.value,
                    "klira.compliance.decision.layer": layer,
                    "klira.guardrails.augmentation_applied": result.augmentation_applied,
                    "klira.guardrails.guidelines_count": len(result.guidelines),
                }
                if matches:
                    attrs["klira.compliance.policy_ids"] = ",".join(
                        m.policy.id for m in matches
                    )
                if result.regulation:
                    attrs["klira.compliance.regulation"] = result.regulation

                # Propagate audit/review flags from matched policies
                audit_required = any(m.policy.audit_required for m in matches)
                human_review = any(m.policy.human_review_required for m in matches)
                attrs["klira.compliance.audit_required"] = audit_required
                attrs["klira.compliance.human_review_required"] = human_review

                with self._tracer.start_as_current_span(span_name, attributes=attrs):
                    pass

            except Exception:
                logger.warning("Failed to create compliance span", exc_info=True)

        _create_span()
