"""
©AngelaMos | 2026
rule.py
"""

from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator

from app.config import settings
from app.models.CorrelationRule import RuleType
from app.models.LogEvent import Severity


class ThresholdConditions(BaseModel):
    """
    Conditions for threshold-based correlation rules
    """
    event_filter: dict[str, str]
    threshold: int = Field(ge=1)
    window_seconds: int = Field(ge=1)
    group_by: str


class SequenceStep(BaseModel):
    """
    A single step in a sequence correlation rule
    """
    event_filter: dict[str, str]
    count: int = Field(default=1, ge=1)


class SequenceConditions(BaseModel):
    """
    Conditions for sequence-based correlation rules
    """
    steps: list[SequenceStep] = Field(min_length=2)
    window_seconds: int = Field(ge=1)
    group_by: str


class AggregationConditions(BaseModel):
    """
    Conditions for aggregation-based correlation rules
    """
    event_filter: dict[str, str]
    aggregation: Literal["count_distinct"]
    aggregation_field: str
    threshold: int = Field(ge=1)
    window_seconds: int = Field(ge=1)
    group_by: str


CONDITION_SCHEMAS: dict[RuleType, type[BaseModel]] = {
    RuleType.THRESHOLD: ThresholdConditions,
    RuleType.SEQUENCE: SequenceConditions,
    RuleType.AGGREGATION: AggregationConditions,
}


class RuleCreateRequest(BaseModel):
    """
    Schema for creating a new correlation rule
    """
    name: str = Field(min_length=1)
    description: str = ""
    rule_type: RuleType
    conditions: dict[str, Any]
    severity: Severity
    enabled: bool = True
    mitre_tactic: str | None = None
    mitre_technique: str | None = None

    @model_validator(mode="after")
    def validate_conditions(self) -> RuleCreateRequest:
        """
        Validate conditions dict matches the expected structure for rule_type
        """
        schema = CONDITION_SCHEMAS[self.rule_type]
        schema.model_validate(self.conditions)
        return self


class RuleUpdateRequest(BaseModel):
    """
    Schema for updating an existing correlation rule
    """
    name: str | None = None
    description: str | None = None
    conditions: dict[str, Any] | None = None
    severity: Severity | None = None
    enabled: bool | None = None
    mitre_tactic: str | None = None
    mitre_technique: str | None = None


class RuleTestRequest(BaseModel):
    """
    Schema for testing a rule against historical events
    """
    hours: int = Field(
        default=settings.TIMELINE_DEFAULT_HOURS,
        ge=1,
        le=settings.RULE_TEST_MAX_HOURS,
    )