"""
ontology.py — Open FAIR taxonomy: enums, assets, threat agents, controls,
frequency/loss factors, scenarios, and risk registers.

The classes in this module form the data model that drives simulations.
They are intentionally lightweight dataclasses so they can be serialised,
copied, and modified without side-effects.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

from .distributions import Distribution


# ---------------------------------------------------------------------------
# Enumerations
# ---------------------------------------------------------------------------

class AssetType(Enum):
    """High-level classification of the asset at risk."""
    INFORMATION = "Information"
    SYSTEM = "System"
    SERVICE = "Service"
    PHYSICAL = "Physical"
    PERSON = "Person"
    OTHER = "Other"


class ThreatType(Enum):
    """Classification of the threat agent's nature."""
    EXTERNAL_ADVERSARIAL = "External Adversarial"
    INTERNAL_ADVERSARIAL = "Internal Adversarial"
    EXTERNAL_NON_ADVERSARIAL = "External Non-Adversarial"
    INTERNAL_NON_ADVERSARIAL = "Internal Non-Adversarial"
    NATURAL = "Natural"


class LossCategory(Enum):
    """Primary and secondary loss categories defined by Open FAIR."""
    PRODUCTIVITY = "Productivity"
    RESPONSE = "Response"
    REPLACEMENT = "Replacement"
    COMPETITIVE_ADVANTAGE = "Competitive Advantage"
    FINES_AND_JUDGMENTS = "Fines and Judgments"
    REPUTATION = "Reputation"


class EffectType(Enum):
    """How the threat agent affects the asset."""
    CONFIDENTIALITY = "Confidentiality"
    INTEGRITY = "Integrity"
    AVAILABILITY = "Availability"
    SAFETY = "Safety"


# ---------------------------------------------------------------------------
# Asset
# ---------------------------------------------------------------------------

@dataclass
class Asset:
    """Represents an organisational asset that may be harmed by a threat.

    Parameters
    ----------
    name:        Human-readable identifier for this asset.
    asset_type:  Classification from :class:`AssetType`.
    description: Optional free-text description.
    tags:        Optional list of string labels for filtering/grouping.
    metadata:    Arbitrary key-value pairs for downstream tooling.
    """

    name: str
    asset_type: AssetType = AssetType.INFORMATION
    description: str = ""
    tags: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("Asset name cannot be empty.")
        if isinstance(self.asset_type, str):
            self.asset_type = AssetType(self.asset_type)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "asset_type": self.asset_type.value,
            "description": self.description,
            "tags": list(self.tags),
            "metadata": dict(self.metadata),
        }


# ---------------------------------------------------------------------------
# ThreatAgent
# ---------------------------------------------------------------------------

@dataclass
class ThreatAgent:
    """Represents an entity capable of causing harm to an asset.

    Parameters
    ----------
    name:        Human-readable identifier.
    threat_type: Classification from :class:`ThreatType`.
    description: Optional free-text description.
    tags:        Optional list of string labels.
    metadata:    Arbitrary key-value pairs.
    """

    name: str
    threat_type: ThreatType = ThreatType.EXTERNAL_ADVERSARIAL
    description: str = ""
    tags: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("ThreatAgent name cannot be empty.")
        if isinstance(self.threat_type, str):
            self.threat_type = ThreatType(self.threat_type)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "threat_type": self.threat_type.value,
            "description": self.description,
            "tags": list(self.tags),
            "metadata": dict(self.metadata),
        }


# ---------------------------------------------------------------------------
# Control
# ---------------------------------------------------------------------------

@dataclass
class Control:
    """Represents a security or operational control that reduces risk.

    Parameters
    ----------
    name:         Human-readable identifier.
    description:  What the control does.
    control_type: Categorical label (e.g., 'Preventive', 'Detective').
    tags:         Optional list of string labels.
    metadata:     Arbitrary key-value pairs.
    """

    name: str
    description: str = ""
    control_type: str = "Preventive"
    tags: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("Control name cannot be empty.")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "control_type": self.control_type,
            "tags": list(self.tags),
            "metadata": dict(self.metadata),
        }


# ---------------------------------------------------------------------------
# FrequencyFactors
# ---------------------------------------------------------------------------

@dataclass
class FrequencyFactors:
    """Captures the distributions used to model how often loss events occur.

    Open FAIR decomposes Threat Event Frequency (TEF) as:
        TEF = Contact Frequency × Probability of Action

    Vulnerability is separately decomposed as:
        Vulnerability = f(Threat Capability, Control Strength)

    Loss Event Frequency (LEF) = TEF × Vulnerability

    Any combination of the fields below is valid; the simulator applies
    a priority ordering to determine which inputs to use.

    Parameters
    ----------
    tef:                   Direct Threat Event Frequency distribution (events/year).
    contact_frequency:     How often the threat agent comes into contact with the asset.
    probability_of_action: Given contact, the probability the agent acts.
    threat_capability:     Threat agent skill/capability, expressed as a 0-100 score distribution.
    control_strength:      Strength of controls, expressed as a 0-100 score distribution.
    vulnerability:         Direct Vulnerability distribution (fraction 0-1).
    """

    tef: Optional[Distribution] = None
    contact_frequency: Optional[Distribution] = None
    probability_of_action: Optional[Distribution] = None
    threat_capability: Optional[Distribution] = None
    control_strength: Optional[Distribution] = None
    vulnerability: Optional[Distribution] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            k: v.to_dict() if v is not None else None
            for k, v in {
                "tef": self.tef,
                "contact_frequency": self.contact_frequency,
                "probability_of_action": self.probability_of_action,
                "threat_capability": self.threat_capability,
                "control_strength": self.control_strength,
                "vulnerability": self.vulnerability,
            }.items()
        }


# ---------------------------------------------------------------------------
# LossMagnitudeFactors
# ---------------------------------------------------------------------------

@dataclass
class LossMagnitudeFactors:
    """Distributions for primary and secondary loss magnitude components.

    Primary loss categories (applied per loss event):
      - productivity:          Revenue/productivity impact.
      - response:              Incident response costs.
      - replacement:           Asset replacement/restoration costs.
      - competitive_advantage: Loss of competitive position.
      - fines_judgments:       Regulatory fines and legal judgments.
      - reputation:            Brand/reputation damage costs.

    A top-level ``primary_loss`` distribution may substitute for all
    individual categories if a single estimate is available.

    Secondary loss is triggered probabilistically:
      - secondary_loss_event_frequency: Probability (0-1) or rate that
        secondary stakeholders (e.g., regulators) act.
      - secondary_loss_magnitude:       Cost per secondary event.
    """

    primary_loss: Optional[Distribution] = None
    productivity: Optional[Distribution] = None
    response: Optional[Distribution] = None
    replacement: Optional[Distribution] = None
    competitive_advantage: Optional[Distribution] = None
    fines_judgments: Optional[Distribution] = None
    reputation: Optional[Distribution] = None
    secondary_loss_event_frequency: Optional[Distribution] = None
    secondary_loss_magnitude: Optional[Distribution] = None

    def to_dict(self) -> Dict[str, Any]:
        fields = [
            "primary_loss", "productivity", "response", "replacement",
            "competitive_advantage", "fines_judgments", "reputation",
            "secondary_loss_event_frequency", "secondary_loss_magnitude",
        ]
        return {f: getattr(self, f).to_dict() if getattr(self, f) is not None else None
                for f in fields}


# ---------------------------------------------------------------------------
# FAIRScenario
# ---------------------------------------------------------------------------

@dataclass
class FAIRScenario:
    """A complete Open FAIR risk scenario ready for Monte Carlo simulation.

    Parameters
    ----------
    name:                  Unique identifier for this scenario.
    asset:                 The asset at risk.
    threat_agent:          The entity posing the threat.
    frequency_factors:     Distributions governing how often events occur.
    loss_magnitude_factors: Distributions governing how much loss per event.
    description:           Optional free-text narrative.
    controls:              List of controls currently in place.
    effect_type:           CIA impact type.
    time_horizon_years:    Simulation time horizon (default 1 year).
    tags:                  Optional labels for grouping/filtering.
    metadata:              Arbitrary key-value pairs.
    """

    name: str
    asset: Asset
    threat_agent: ThreatAgent
    frequency_factors: FrequencyFactors
    loss_magnitude_factors: LossMagnitudeFactors
    description: str = ""
    controls: List[Control] = field(default_factory=list)
    effect_type: EffectType = EffectType.CONFIDENTIALITY
    time_horizon_years: float = 1.0
    tags: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("FAIRScenario name cannot be empty.")
        if self.time_horizon_years <= 0:
            raise ValueError(
                f"time_horizon_years must be positive; got {self.time_horizon_years}"
            )
        if isinstance(self.effect_type, str):
            self.effect_type = EffectType(self.effect_type)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "asset": self.asset.to_dict(),
            "threat_agent": self.threat_agent.to_dict(),
            "frequency_factors": self.frequency_factors.to_dict(),
            "loss_magnitude_factors": self.loss_magnitude_factors.to_dict(),
            "description": self.description,
            "controls": [c.to_dict() for c in self.controls],
            "effect_type": self.effect_type.value,
            "time_horizon_years": self.time_horizon_years,
            "tags": list(self.tags),
            "metadata": dict(self.metadata),
        }


# ---------------------------------------------------------------------------
# RiskRegister
# ---------------------------------------------------------------------------

@dataclass
class RiskRegister:
    """An ordered collection of :class:`FAIRScenario` objects.

    Parameters
    ----------
    name:        Human-readable identifier for this register.
    scenarios:   Initial list of scenarios (optional).
    description: Free-text description of the register's scope.
    metadata:    Arbitrary key-value pairs.
    """

    name: str
    scenarios: List[FAIRScenario] = field(default_factory=list)
    description: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("RiskRegister name cannot be empty.")
        # Convert list to a mutable copy to avoid aliasing issues
        self.scenarios = list(self.scenarios)

    # -- Mutation helpers --

    def add_scenario(self, scenario: FAIRScenario) -> "RiskRegister":
        """Append *scenario* to the register.  Returns self for chaining."""
        if not isinstance(scenario, FAIRScenario):
            raise TypeError(f"Expected FAIRScenario, got {type(scenario).__name__}")
        existing_names = [s.name for s in self.scenarios]
        if scenario.name in existing_names:
            raise ValueError(
                f"A scenario named '{scenario.name}' already exists in this register."
            )
        self.scenarios.append(scenario)
        return self

    def remove_scenario(self, name: str) -> "RiskRegister":
        """Remove the scenario with the given *name*.  Returns self for chaining."""
        original_len = len(self.scenarios)
        self.scenarios = [s for s in self.scenarios if s.name != name]
        if len(self.scenarios) == original_len:
            raise KeyError(f"No scenario named '{name}' found in this register.")
        return self

    def get_scenario(self, name: str) -> FAIRScenario:
        """Return the scenario with the given *name*, or raise :class:`KeyError`."""
        for scenario in self.scenarios:
            if scenario.name == name:
                return scenario
        raise KeyError(f"No scenario named '{name}' found in this register.")

    def __len__(self) -> int:
        return len(self.scenarios)

    def __iter__(self):
        return iter(self.scenarios)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "scenarios": [s.to_dict() for s in self.scenarios],
            "metadata": dict(self.metadata),
        }
