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

"""Guardrails lifecycle contract — state machine specification for guardrails processing.

Defines the expected sequence of operations, span creation, and data flow
for guardrails evaluation. This contract is validated by dedicated tests in
tests/contracts/test_guardrails_lifecycle.py — it is not enforced at runtime
by the engine.

States:
    IDLE → EVALUATING → DECIDED → AUDIT_SCHEDULED → DONE

Key invariants:
    - Guidelines pass via explicit GuardrailProcessingResult.guidelines,
      NEVER via OTel context. (Learning #7)
    - augmentation_applied set ONLY after verifying injection. (Learning #7)
    - Compliance audit spans are always async — never block hot path. (Learning #20)
    - Two policy actions only: block or allow. No "warn".
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum


class GuardrailState(Enum):
    """States in the guardrails processing lifecycle."""

    IDLE = "idle"
    EVALUATING = "evaluating"
    DECIDED = "decided"
    AUDIT_SCHEDULED = "audit_scheduled"
    DONE = "done"


@dataclass(frozen=True)
class GuardrailTransition:
    """A valid state transition in the guardrails lifecycle."""

    from_state: GuardrailState
    to_state: GuardrailState
    description: str
    span_created: str | None = None
    is_async: bool = False


VALID_TRANSITIONS: tuple[GuardrailTransition, ...] = (
    GuardrailTransition(
        from_state=GuardrailState.IDLE,
        to_state=GuardrailState.EVALUATING,
        description=(
            "FastRulesEngine starts evaluation. Creates guardrails.input or "
            "guardrails.output span."
        ),
        span_created="klira.guardrails.input",
    ),
    GuardrailTransition(
        from_state=GuardrailState.EVALUATING,
        to_state=GuardrailState.DECIDED,
        description=(
            "DecisionRouter determines action from FastRulesEngine results. "
            "Sets decision/allowed attributes on parent span. "
            "Returns GuardrailProcessingResult."
        ),
    ),
    GuardrailTransition(
        from_state=GuardrailState.DECIDED,
        to_state=GuardrailState.AUDIT_SCHEDULED,
        description=(
            "Compliance audit span scheduled asynchronously. "
            "Never blocks the hot path. Span type depends on decision: "
            "klira.compliance.{allowed|blocked|augmented|llm_fallback}."
        ),
        span_created="klira.compliance.*",
        is_async=True,
    ),
    GuardrailTransition(
        from_state=GuardrailState.AUDIT_SCHEDULED,
        to_state=GuardrailState.DONE,
        description="Result returned to caller. Lifecycle complete.",
    ),
)


class GuardrailLifecycle:
    """Tracks and validates guardrails lifecycle state transitions.

    Usage:
        lifecycle = GuardrailLifecycle()
        lifecycle.transition_to(GuardrailState.EVALUATING)
        lifecycle.transition_to(GuardrailState.DECIDED)
        lifecycle.transition_to(GuardrailState.AUDIT_SCHEDULED)
        lifecycle.transition_to(GuardrailState.DONE)
    """

    _valid_transitions: dict[tuple[GuardrailState, GuardrailState], GuardrailTransition]

    def __init__(self) -> None:
        self._state = GuardrailState.IDLE
        self._valid_transitions = {
            (t.from_state, t.to_state): t for t in VALID_TRANSITIONS
        }

    @property
    def state(self) -> GuardrailState:
        """Current lifecycle state."""
        return self._state

    def transition_to(self, new_state: GuardrailState) -> GuardrailTransition:
        """Transition to a new state.

        Args:
            new_state: The target state.

        Returns:
            The transition that was applied.

        Raises:
            ValueError: If the transition is not valid from the current state.
        """
        key = (self._state, new_state)
        transition = self._valid_transitions.get(key)
        if transition is None:
            valid_targets = [
                t.to_state.value
                for t in VALID_TRANSITIONS
                if t.from_state == self._state
            ]
            raise ValueError(
                f"Invalid transition: {self._state.value} → {new_state.value}. "
                f"Valid transitions from {self._state.value}: {valid_targets}"
            )
        self._state = new_state
        return transition

    def reset(self) -> None:
        """Reset the lifecycle to IDLE state."""
        self._state = GuardrailState.IDLE

    @property
    def is_complete(self) -> bool:
        """Whether the lifecycle has reached DONE state."""
        return self._state == GuardrailState.DONE
