# Copyright (c) 2025, HaiyangLi <quantocean.li at gmail dot com>
# SPDX-License-Identifier: Apache-2.0

"""Base rule system for validation and auto-correction.

Resurrected from lionagi v0.2.2 and refined from lionherd-old patterns.
This is the validation layer of the IPU (Intelligence Processing Unit).
"""

from abc import abstractmethod
from dataclasses import dataclass, field
from enum import IntEnum, auto
from typing import Any

from lionpride.errors import ValidationError
from lionpride.types import Params

__all__ = ("Rule", "RuleParams", "RuleQualifier", "ValidationError")


class RuleQualifier(IntEnum):
    """Qualifier types for rules - determines WHEN a rule applies.

    - FIELD: Match by field name (e.g., "confidence", "output")
    - ANNOTATION: Match by type annotation (e.g., str, int, float)
    - CONDITION: Match by custom condition (e.g., is BaseModel subclass)

    Default precedence order: FIELD > ANNOTATION > CONDITION
    """

    FIELD = auto()
    ANNOTATION = auto()
    CONDITION = auto()

    @classmethod
    def from_str(cls, s: str) -> "RuleQualifier":
        """Convert string to RuleQualifier."""
        s = s.strip().upper()
        if s == "FIELD":
            return cls.FIELD
        elif s == "ANNOTATION":
            return cls.ANNOTATION
        elif s == "CONDITION":
            return cls.CONDITION
        else:
            raise ValueError(f"Unknown RuleQualifier: {s}")


def _decide_qualifier_order(
    qualifier: str | RuleQualifier | None = None,
) -> list[RuleQualifier]:
    """Determine qualifier precedence order.

    Args:
        qualifier: Optional preferred qualifier (tried first)

    Returns:
        List of qualifiers in precedence order
    """
    default_order = [
        RuleQualifier.FIELD,
        RuleQualifier.ANNOTATION,
        RuleQualifier.CONDITION,
    ]

    if qualifier is None:
        return default_order

    if isinstance(qualifier, str):
        qualifier = RuleQualifier.from_str(qualifier)

    # Move specified qualifier to front
    default_order.remove(qualifier)
    return [qualifier, *default_order]


@dataclass(slots=True, frozen=True)
class RuleParams(Params):
    """Immutable configuration for rules.

    Defines:
    - WHAT the rule applies to (types or fields)
    - HOW to determine applicability (default_qualifier)
    - WHETHER to auto-fix (auto_fix)
    - ADDITIONAL validation parameters (kw)

    Uses lionpride.types.Params (leaner than Pydantic BaseModel).
    """

    apply_types: set[type] = field(default_factory=set)
    """Types this rule applies to (e.g., {str, int})"""

    apply_fields: set[str] = field(default_factory=set)
    """Field names this rule applies to (e.g., {"confidence", "output"})"""

    default_qualifier: RuleQualifier = RuleQualifier.FIELD
    """Preferred qualifier type"""

    auto_fix: bool = False
    """Enable automatic fixing on validation failure"""

    kw: dict = field(default_factory=dict)
    """Additional validation parameters"""

    def _validate(self) -> None:
        """Validate params - ensure exactly one of apply_types or apply_fields is set."""
        super()._validate()

        both_set = bool(self.apply_types) and bool(self.apply_fields)
        neither_set = not self.apply_types and not self.apply_fields

        if both_set:
            raise ValueError("Cannot set both apply_types and apply_fields")
        if neither_set and self.default_qualifier != RuleQualifier.CONDITION:
            raise ValueError(
                "Must set apply_types or apply_fields unless using CONDITION qualifier"
            )


class Rule:
    """Base validation rule with auto-correction support.

    Pattern from lionagi v0.2.2 + lionherd-old:
    1. apply() - Does this rule apply? (uses qualifiers)
    2. validate() - Is the value valid? (abstract, subclass implements)
    3. perform_fix() - Can we auto-correct? (optional, if auto_fix=True)

    Usage:
        # Create rule
        rule = StringRule(min_length=1, max_length=100)

        # Check if applies
        if await rule.apply("name", "Ocean", str):
            # Validate + fix
            result = await rule.invoke("name", "Ocean", str)
    """

    def __init__(self, params: RuleParams, **kw):
        """Initialize rule with parameters.

        Args:
            params: Rule configuration
            **kw: Additional validation kwargs (merged with params.kw)
        """
        if kw:
            # Merge additional kwargs using with_updates
            params = params.with_updates(kw={**params.kw, **kw})
        self.params = params

    @property
    def apply_types(self) -> set[type]:
        """Types this rule applies to."""
        return self.params.apply_types

    @property
    def apply_fields(self) -> set[str]:
        """Field names this rule applies to."""
        return self.params.apply_fields

    @property
    def default_qualifier(self) -> RuleQualifier:
        """Preferred qualifier type."""
        return self.params.default_qualifier

    @property
    def auto_fix(self) -> bool:
        """Whether auto-fixing is enabled."""
        return self.params.auto_fix

    @property
    def validation_kwargs(self) -> dict:
        """Additional validation parameters."""
        return self.params.kw

    async def rule_condition(self, k: str, v: Any, t: type, **kw) -> bool:
        """Custom condition for CONDITION qualifier.

        Override in subclass to use CONDITION qualifier.

        Args:
            k: Field name
            v: Field value
            t: Field type
            **kw: Additional kwargs

        Returns:
            True if rule should apply
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} must implement rule_condition() to use CONDITION qualifier"
        )

    async def _apply(self, k: str, v: Any, t: type, q: RuleQualifier, **kw) -> bool:
        """Determine if rule applies based on qualifier.

        Args:
            k: Field name
            v: Field value
            t: Field type
            q: Qualifier type
            **kw: Additional kwargs

        Returns:
            True if rule applies
        """
        match q:
            case RuleQualifier.FIELD:
                return k in self.apply_fields

            case RuleQualifier.ANNOTATION:
                return t in self.apply_types

            case RuleQualifier.CONDITION:
                return await self.rule_condition(k, v, t, **kw)

    async def apply(
        self,
        k: str,
        v: Any,
        t: type | None = None,
        qualifier: str | RuleQualifier | None = None,
        **kw,
    ) -> bool:
        """Check if rule applies using qualifier precedence.

        Args:
            k: Field name
            v: Field value
            t: Field type (optional)
            qualifier: Override default qualifier order
            **kw: Additional kwargs for condition checking

        Returns:
            True if rule applies (any qualifier matches)
        """
        _order = _decide_qualifier_order(qualifier)

        for q in _order:
            try:
                if await self._apply(k, v, t or type(v), q, **kw):
                    return True
            except NotImplementedError:
                # Skip qualifiers not implemented by this rule
                continue

        return False

    @abstractmethod
    async def validate(self, v: Any, t: type, **kw) -> None:
        """Validate value (abstract, implement in subclass).

        Args:
            v: Value to validate
            t: Expected type
            **kw: Additional validation parameters

        Raises:
            Exception: If validation fails
        """
        pass

    async def perform_fix(self, v: Any, t: type) -> Any:
        """Attempt to fix invalid value (optional, override in subclass).

        Args:
            v: Value to fix
            t: Expected type

        Returns:
            Fixed value

        Raises:
            NotImplementedError: If auto_fix=True but perform_fix not implemented
        """
        raise NotImplementedError(
            f"{self.__class__.__name__} must implement perform_fix() to use auto_fix=True"
        )

    async def invoke(
        self, k: str, v: Any, t: type | None = None, *, auto_fix: bool | None = None
    ) -> Any:
        """Execute validation with optional auto-fixing.

        Args:
            k: Field name (for error messages)
            v: Value to validate
            t: Field type (optional)
            auto_fix: Override self.auto_fix for this invocation (thread-safe)

        Returns:
            Validated (and possibly fixed) value

        Raises:
            ValidationError: If validation fails and auto_fix disabled
        """
        effective_type = t or type(v)
        # Use override if provided, otherwise fall back to self.auto_fix
        should_auto_fix = auto_fix if auto_fix is not None else self.auto_fix
        try:
            await self.validate(v, effective_type, **self.validation_kwargs)
            return v
        except Exception as e:
            if should_auto_fix:
                try:
                    return await self.perform_fix(v, effective_type)
                except Exception as e1:
                    raise ValidationError(f"Failed to fix field '{k}': {e1}") from e
            raise ValidationError(f"Failed to validate field '{k}': {e}") from e

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"{self.__class__.__name__}("
            f"types={self.apply_types}, "
            f"fields={self.apply_fields}, "
            f"auto_fix={self.auto_fix})"
        )
