# scm/models/security/vulnerability_protection_profiles.py

from typing import List, Optional
from uuid import UUID

from pydantic import (
    BaseModel,
    Field,
    model_validator,
    ConfigDict,
    RootModel,
)
from enum import Enum


# Keep all original Enums
class Severity(str, Enum):
    """Enumeration of severity levels."""

    critical = "critical"
    high = "high"
    medium = "medium"
    low = "low"
    informational = "informational"
    any = "any"


class Category(str, Enum):
    """Enumeration of vulnerability categories."""

    any = "any"
    brute_force = "brute-force"
    code_execution = "code-execution"
    code_obfuscation = "code-obfuscation"
    command_execution = "command-execution"
    dos = "dos"
    exploit_kit = "exploit-kit"
    info_leak = "info-leak"
    insecure_credentials = "insecure-credentials"
    overflow = "overflow"
    phishing = "phishing"
    protocol_anomaly = "protocol-anomaly"
    scan = "scan"
    sql_injection = "sql-injection"


class PacketCapture(str, Enum):
    """Enumeration of packet capture options."""

    disable = "disable"
    single_packet = "single-packet"
    extended_capture = "extended-capture"


class Host(str, Enum):
    """Enumeration of host options."""

    any = "any"
    client = "client"
    server = "server"


class BlockIpTrackBy(str, Enum):
    """Enumeration for block_ip track_by options."""

    source_and_destination = "source-and-destination"
    source = "source"


class TimeAttributeTrackBy(str, Enum):
    """Enumeration for time_attribute track_by options."""

    source = "source"
    destination = "destination"
    source_and_destination = "source-and-destination"


# Component Models (keep all original component models)
class BlockIpAction(BaseModel):
    """Represents the 'block_ip' action with additional properties."""

    track_by: BlockIpTrackBy = Field(..., description="Tracking method")
    duration: int = Field(..., description="Duration in seconds", ge=1, le=3600)


class ActionRequest(RootModel[dict]):
    """
    Represents the 'action' field in rules and threat exceptions for requests.
    Enforces that exactly one action is provided and validates 'block_ip' properties.
    """

    @model_validator(mode="before")
    def check_and_transform_action(cls, values):
        if isinstance(values, str):
            # Convert string to dict
            values = {values: {}}
        elif isinstance(values, dict):
            pass
        else:
            raise ValueError("Invalid action format; must be a string or dict.")

        action_fields = [
            "allow",
            "alert",
            "drop",
            "reset_client",
            "reset_server",
            "reset_both",
            "block_ip",
            "default",
        ]

        provided_actions = [field for field in action_fields if field in values]

        if len(provided_actions) != 1:
            raise ValueError("Exactly one action must be provided in 'action' field.")

        action_name = provided_actions[0]
        action_value = values[action_name]

        if action_name == "block_ip":
            # Validate block_ip properties
            if not isinstance(action_value, dict):
                raise ValueError(
                    "'block_ip' action must be a dict with 'track_by' and 'duration' fields."
                )

            try:
                BlockIpAction(**action_value)
            except Exception as e:
                raise ValueError(f"Invalid 'block_ip' action: {e}")
        else:
            # Other actions should have empty dict as value
            if action_value != {}:
                raise ValueError(
                    f"Action '{action_name}' does not take any parameters."
                )

        return values

    def get_action_name(self) -> str:
        if self.root is None:
            return "unknown"
        return next(iter(self.root.keys()), "unknown")


class ActionResponse(RootModel[dict]):
    """
    Represents the 'action' field in rules and threat exceptions for responses.
    Accepts empty dictionaries.
    """

    @model_validator(mode="before")
    def check_action(cls, values):
        if isinstance(values, str):
            values = {values: {}}
        elif isinstance(values, dict):
            pass
        else:
            raise ValueError("Invalid action format; must be a string or dict.")

        action_fields = [
            "allow",
            "alert",
            "drop",
            "reset_client",
            "reset_server",
            "reset_both",
            "block_ip",
            "default",
        ]

        provided_actions = [field for field in action_fields if field in values]

        if len(provided_actions) > 1:
            raise ValueError("At most one action must be provided in 'action' field.")

        if provided_actions:
            action_name = provided_actions[0]
            action_value = values[action_name]

            if action_name == "block_ip":
                # Validate block_ip properties
                if not isinstance(action_value, dict):
                    raise ValueError(
                        "'block_ip' action must be a dict with 'track_by' and 'duration' fields."
                    )

                try:
                    BlockIpAction(**action_value)
                except Exception as e:
                    raise ValueError(f"Invalid 'block_ip' action: {e}")
            else:
                # Other actions should have empty dict as value
                if action_value != {}:
                    raise ValueError(
                        f"Action '{action_name}' does not take any parameters."
                    )
        else:
            # Accept empty dicts (no action specified)
            if values != {}:
                raise ValueError("Invalid action format.")

        return values

    def get_action_name(self) -> str:
        if self.root is None:
            return "unknown"
        return next(iter(self.root.keys()), "unknown")


class TimeAttribute(BaseModel):
    """Represents the 'time_attribute' field in threat exceptions."""

    interval: int = Field(
        ...,
        description="Interval in seconds",
        ge=1,
        le=3600,
    )
    threshold: int = Field(
        ...,
        description="Threshold",
        ge=1,
        le=65535,
    )
    track_by: TimeAttributeTrackBy = Field(
        ...,
        description="Track by",
    )


class ExemptIpEntry(BaseModel):
    """Represents an entry in the 'exempt_ip' list."""

    name: str = Field(..., description="Exempt IP name")


# Rule Models
class VulnerabilityRuleModel(BaseModel):
    """Base model for rules containing common fields."""

    name: str = Field(
        ...,
        description="Rule name",
    )
    packet_capture: PacketCapture = Field(
        default="disable",
        description="Packet capture setting",
    )
    severity: List[Severity] = Field(
        ...,
        description="List of severities",
    )
    category: Category = Field(
        default="any",
        description="Category",
    )
    cve: List[str] = Field(
        default=["any"],
        description="List of CVEs",
    )
    host: Host = Field(
        ...,
        description="Host",
    )
    vendor_id: List[str] = Field(
        default=["any"],
        description="List of vendor IDs",
        alias="vendor-id",
    )
    threat_name: str = Field(
        default="any",
        description="Threat name",
    )
    action: Optional[ActionResponse] = Field(
        None,
        description="Action for the rule",
    )

    model_config = ConfigDict(populate_by_name=True, alias_generator=None)


# Threat Exception Models
class ThreatExceptionModel(BaseModel):
    """Base model for threat exceptions containing common fields."""

    name: str = Field(
        ...,
        description="Threat exception name",
    )
    packet_capture: Optional[PacketCapture] = Field(
        None,
        description="Packet capture setting",
    )
    exempt_ip: Optional[List[ExemptIpEntry]] = Field(
        None,
        description="Exempt IP list",
    )
    time_attribute: Optional[TimeAttribute] = Field(
        None,
        description="Time attribute",
    )
    notes: Optional[str] = Field(
        None,
        description="Notes",
    )


# Main Models
class VulnerabilityProtectionProfileBaseModel(BaseModel):
    """Base model containing common fields for all Vulnerability Protection Profile operations."""

    model_config = ConfigDict(
        validate_assignment=True,
        arbitrary_types_allowed=True,
        populate_by_name=True,
    )

    name: str = Field(
        ...,
        description="Profile name",
        pattern=r"^[a-zA-Z0-9._-]+$",
    )
    rules: List[VulnerabilityRuleModel] = Field(
        ...,
        description="List of rules",
    )
    threat_exception: Optional[List[ThreatExceptionModel]] = Field(
        None,
        description="List of threat exceptions",
    )
    description: Optional[str] = Field(
        None,
        description="Description",
    )
    folder: Optional[str] = Field(
        None,
        description="Folder",
        max_length=64,
        pattern=r"^[a-zA-Z\d\-_. ]+$",
    )
    snippet: Optional[str] = Field(
        None,
        description="Snippet",
        max_length=64,
        pattern=r"^[a-zA-Z\d\-_. ]+$",
    )
    device: Optional[str] = Field(
        None,
        description="Device",
        max_length=64,
        pattern=r"^[a-zA-Z\d\-_. ]+$",
    )


class VulnerabilityProtectionProfileCreateModel(
    VulnerabilityProtectionProfileBaseModel
):
    """Model for creating a new Vulnerability Protection Profile."""

    @model_validator(mode="after")
    def validate_container(self) -> "VulnerabilityProtectionProfileCreateModel":
        container_fields = [
            "folder",
            "snippet",
            "device",
        ]
        provided = [
            field for field in container_fields if getattr(self, field) is not None
        ]
        if len(provided) != 1:
            raise ValueError(
                "Exactly one of 'folder', 'snippet', or 'device' must be provided."
            )
        return self


class VulnerabilityProtectionProfileUpdateModel(
    VulnerabilityProtectionProfileBaseModel
):
    """Model for updating an existing Vulnerability Protection Profile."""


class VulnerabilityProtectionProfileResponseModel(
    VulnerabilityProtectionProfileBaseModel
):
    """Model for Vulnerability Protection Profile responses."""

    id: UUID = Field(
        ...,
        description="Profile ID",
        examples=["123e4567-e89b-12d3-a456-426655440000"],
    )
