/tests/synthetic/schemas.py

"""
Centralized schema definitions for synthetic testing fixtures.

All scenario fixture files (alert.json, aws_cloudwatch_metrics.json, aws_rds_events.json,
aws_performance_insights.json, answer.yml, scenario.yml) must conform to these TypedDicts.
Validators enforce required fields so every scenario is structurally consistent.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, NotRequired

from typing_extensions import TypedDict

# ---------------------------------------------------------------------------
# Controlled vocabularies for scenario metadata
# ---------------------------------------------------------------------------

VALID_ENGINES = frozenset({"postgres", "mysql", "aurora-postgres", "aurora-mysql", "mariadb"})
VALID_FAILURE_MODES = frozenset(
    {
        "replication_lag",
        "connection_exhaustion",
        "storage_full",
        "cpu_saturation",
        "failover",
        "healthy",
        "application_load_spike",
    }
)
VALID_EVIDENCE_SOURCES = frozenset(
    {
        "aws_cloudwatch_metrics",
        "aws_rds_events",
        "aws_performance_insights",
        "ec2_instances_by_tag",
        "elb_target_health",
        "k8s_events",
        "k8s_pod_metrics",
        "k8s_node_metrics",
        "k8s_dns_metrics",
        "k8s_mesh_metrics",
        "k8s_rollout",
    }
)

# ---------------------------------------------------------------------------
# Alert fixture  (alert.json)
# ---------------------------------------------------------------------------


class AlertLabels(TypedDict, total=False):
    alertname: str
    severity: str
    pipeline_name: str
    service: str
    engine: str


class AlertAnnotations(TypedDict, total=False):
    summary: str
    error: str
    suspected_symptom: str
    db_instance_identifier: str
    db_instance: str
    db_cluster: str
    read_replica: str
    cloudwatch_region: str
    rds_failure_mode: str
    context_sources: str


class AlertFixture(TypedDict):
    title: str
    state: str
    alert_source: str
    commonLabels: AlertLabels
    commonAnnotations: AlertAnnotations


# ---------------------------------------------------------------------------
# CloudWatch metrics fixture  (aws_cloudwatch_metrics.json)
# Models the AWS GetMetricData API response shape.
# ---------------------------------------------------------------------------


class MetricDimension(TypedDict):
    Name: str
    Value: str


class MetricDataResult(TypedDict):
    """One metric query result, combining query context with response data."""

    id: str
    label: str
    metric_name: str
    dimensions: list[MetricDimension]
    stat: str
    unit: str
    status_code: str
    timestamps: list[str]
    values: list[float]


class CloudWatchMetricsFixture(TypedDict):
    namespace: str
    period: int
    start_time: str
    end_time: str
    metric_data_results: list[MetricDataResult]


# ---------------------------------------------------------------------------
# RDS events fixture  (aws_rds_events.json)
# Models the AWS DescribeEvents API response shape.
# ---------------------------------------------------------------------------


class RDSEvent(TypedDict):
    date: str
    message: str
    source_identifier: str
    source_type: str
    event_categories: list[str]


class RDSEventsFixture(TypedDict):
    events: list[RDSEvent]


# ---------------------------------------------------------------------------
# Performance insights fixture  (aws_performance_insights.json)
# Models the AWS GetResourceMetrics + DescribeDimensionKeys API response shape.
# ---------------------------------------------------------------------------


class DBLoadTimeSeries(TypedDict):
    timestamps: list[str]
    values: list[float]
    unit: str


class TopSQLWaitEvent(TypedDict):
    name: str
    type: str
    db_load_avg: float


class TopSQL(TypedDict):
    statement: str
    db_load_avg: float
    wait_events: list[TopSQLWaitEvent]
    calls_per_sec: float


class TopWaitEvent(TypedDict):
    name: str
    type: str
    db_load_avg: float


class TopUser(TypedDict):
    name: str
    db_load_avg: float


class TopHost(TypedDict):
    id: str
    db_load_avg: float


class PerformanceInsightsFixture(TypedDict):
    db_instance_identifier: str
    start_time: str
    end_time: str
    db_load: DBLoadTimeSeries
    top_sql: list[TopSQL]
    top_wait_events: list[TopWaitEvent]
    top_users: list[TopUser]
    top_hosts: list[TopHost]


# ---------------------------------------------------------------------------
# Answer key  (answer.yml)
# ---------------------------------------------------------------------------


VALID_TRAJECTORY_ACTIONS = frozenset(
    {
        "query_grafana_metrics",
        "query_grafana_logs",
        "query_grafana_alert_rules",
        "describe_rds_instance",
        "describe_rds_events",
        "ec2_instances_by_tag",
        "get_elb_target_health",
    }
)


class AnswerKeySchema(TypedDict):
    root_cause_category: str
    required_keywords: list[str]
    model_response: str
    # Additional categories that pass the category gate alongside root_cause_category
    equivalent_root_cause_categories: NotRequired[list[str]]
    # Optional adversarial constraints (level 2+ scenarios)
    forbidden_categories: NotRequired[list[str]]  # root_cause_category must NOT be any of these
    forbidden_keywords: NotRequired[list[str]]  # none of these may appear in evidence_text
    required_evidence_sources: NotRequired[
        list[str]
    ]  # these keys must be non-empty in final_state["evidence"]
    # Trajectory efficiency (Axis 1)
    optimal_trajectory: NotRequired[list[str]]  # ordered action names the agent should call
    max_investigation_loops: NotRequired[int]  # how many investigation loops is acceptable
    # Adversarial reasoning (Axis 2)
    ruling_out_keywords: NotRequired[
        list[str]
    ]  # agent output must contain these tokens (proof it dismissed alternatives)
    required_queries: NotRequired[
        list[str]
    ]  # metric names agent must have specifically requested via query_timeseries
    golden_trajectory: NotRequired[GoldenTrajectorySchema]


class GoldenTrajectorySchema(TypedDict, total=False):
    ordered_actions: list[str]
    matching: str
    max_edit_distance: int
    max_extra_actions: int
    max_redundancy: int
    max_loops: int


# ---------------------------------------------------------------------------
# Scenario metadata  (scenario.yml)
# ---------------------------------------------------------------------------


class TopologyTier(TypedDict):
    """Application tier description for EC2/RDS topology fixtures."""

    name: str
    instance_ids: list[str]
    asg: NotRequired[str]


class TopologyMetadata(TypedDict, total=False):
    """Optional EC2/RDS topology block for non-K8s scenarios.

    Present in fixtures that exercise the DNS → LB → Target Group → EC2 → RDS
    request path. Absent in legacy RDS-only scenarios (000–014).
    """

    vpc_id: str
    load_balancer_arn: str
    target_group_arn: str
    tiers: list[TopologyTier]


class ScenarioMetadataSchema(TypedDict):
    schema_version: str
    scenario_id: str
    engine: str
    engine_version: str
    instance_class: str
    region: str
    db_instance_identifier: str
    failure_mode: str
    severity: str
    available_evidence: list[str]
    db_cluster: NotRequired[str]
    scenario_difficulty: NotRequired[int]  # 1–4 curriculum level
    adversarial_signals: NotRequired[list[str]]  # metrics that are intentional confounders
    depends_on: NotRequired[str]  # e.g. "healthy_rca_state" — CI skip flag
    topology: NotRequired[TopologyMetadata]


# ---------------------------------------------------------------------------
# Typed evidence container
# ---------------------------------------------------------------------------


class EC2Instance(TypedDict, total=False):
    instance_id: str
    tier: str
    asg: str
    private_ip: str
    vpc_id: str
    subnet_id: str
    state: str
    instance_type: str
    security_groups: list[str]


class EC2InstancesByTagFixture(TypedDict):
    instances: list[EC2Instance]


class ELBTargetHealthEntry(TypedDict, total=False):
    target_group_arn: str
    instance_id: str
    port: int
    state: str
    reason: str
    description: str


class ELBTargetGroup(TypedDict, total=False):
    TargetGroupArn: str
    TargetGroupName: str
    LoadBalancerArns: list[str]


class ELBTargetHealthFixture(TypedDict):
    target_groups: list[ELBTargetGroup]
    targets: list[ELBTargetHealthEntry]


class GenericEvidenceFixture(TypedDict, total=False):
    """Flexible schema for suite-specific evidence not strongly typed here."""


@dataclass(frozen=True)
class ScenarioEvidence:
    """Typed container for all evidence sources in a scenario fixture.

    Each attribute is None when the corresponding file was not listed in
    scenario.yml:available_evidence, making evidence presence explicit.
    """

    aws_cloudwatch_metrics: CloudWatchMetricsFixture | None
    aws_rds_events: list[RDSEvent] | None
    aws_performance_insights: PerformanceInsightsFixture | None
    ec2_instances_by_tag: EC2InstancesByTagFixture | None = None
    elb_target_health: ELBTargetHealthFixture | None = None
    k8s_events: GenericEvidenceFixture | None = None
    k8s_pod_metrics: GenericEvidenceFixture | None = None
    k8s_node_metrics: GenericEvidenceFixture | None = None
    k8s_dns_metrics: GenericEvidenceFixture | None = None
    k8s_mesh_metrics: GenericEvidenceFixture | None = None
    k8s_rollout: GenericEvidenceFixture | None = None

    def as_dict(self) -> dict[str, Any]:
        """Return only the non-None sources as a plain dict."""
        result: dict[str, Any] = {}
        if self.aws_cloudwatch_metrics is not None:
            result["aws_cloudwatch_metrics"] = self.aws_cloudwatch_metrics
        if self.aws_rds_events is not None:
            result["aws_rds_events"] = self.aws_rds_events
        if self.aws_performance_insights is not None:
            result["aws_performance_insights"] = self.aws_performance_insights
        if self.ec2_instances_by_tag is not None:
            result["ec2_instances_by_tag"] = self.ec2_instances_by_tag
        if self.elb_target_health is not None:
            result["elb_target_health"] = self.elb_target_health
        if self.k8s_events is not None:
            result["k8s_events"] = self.k8s_events
        if self.k8s_pod_metrics is not None:
            result["k8s_pod_metrics"] = self.k8s_pod_metrics
        if self.k8s_node_metrics is not None:
            result["k8s_node_metrics"] = self.k8s_node_metrics
        if self.k8s_dns_metrics is not None:
            result["k8s_dns_metrics"] = self.k8s_dns_metrics
        if self.k8s_mesh_metrics is not None:
            result["k8s_mesh_metrics"] = self.k8s_mesh_metrics
        if self.k8s_rollout is not None:
            result["k8s_rollout"] = self.k8s_rollout
        return result

    def get(self, key: str) -> Any:
        return self.as_dict().get(key)


# ---------------------------------------------------------------------------
# Validators — raise ValueError with a descriptive message on bad data
# ---------------------------------------------------------------------------


def validate_alert(data: dict[str, Any]) -> AlertFixture:
    _require_str(data, "title", ctx="alert.json")
    _require_str(data, "state", ctx="alert.json")
    _require_str(data, "alert_source", ctx="alert.json")
    if not isinstance(data.get("commonLabels"), dict):
        raise ValueError("alert.json: 'commonLabels' must be an object")
    if not isinstance(data.get("commonAnnotations"), dict):
        raise ValueError("alert.json: 'commonAnnotations' must be an object")
    return data  # type: ignore[return-value]


def validate_cloudwatch_metrics(data: dict[str, Any]) -> CloudWatchMetricsFixture:
    ctx = "aws_cloudwatch_metrics.json"
    _require_str(data, "namespace", ctx=ctx)
    _require_str(data, "start_time", ctx=ctx)
    _require_str(data, "end_time", ctx=ctx)
    if not isinstance(data.get("period"), int):
        raise ValueError(f"{ctx}: 'period' must be an integer (seconds)")
    results = data.get("metric_data_results")
    if not isinstance(results, list) or not results:
        raise ValueError(f"{ctx}: 'metric_data_results' must be a non-empty list")
    for i, result in enumerate(results):
        rctx = f"{ctx}:metric_data_results[{i}]"
        for field in ("id", "label", "metric_name", "stat", "unit", "status_code"):
            _require_str(result, field, ctx=rctx)
        if not isinstance(result.get("dimensions"), list):
            raise ValueError(f"{rctx}: 'dimensions' must be a list")
        for dim in result["dimensions"]:
            _require_str(dim, "Name", ctx=rctx)
            _require_str(dim, "Value", ctx=rctx)
        if not isinstance(result.get("timestamps"), list):
            raise ValueError(f"{rctx}: 'timestamps' must be a list")
        if not isinstance(result.get("values"), list):
            raise ValueError(f"{rctx}: 'values' must be a list")
        if len(result["timestamps"]) != len(result["values"]):
            raise ValueError(f"{rctx}: 'timestamps' and 'values' must have the same length")
    return data  # type: ignore[return-value]


def validate_rds_events(data: dict[str, Any]) -> RDSEventsFixture:
    if not isinstance(data.get("events"), list):
        raise ValueError("aws_rds_events.json: 'events' must be a list")
    for i, event in enumerate(data["events"]):
        ctx = f"aws_rds_events.json:events[{i}]"
        _require_str(event, "date", ctx=ctx)
        _require_str(event, "message", ctx=ctx)
        _require_str(event, "source_identifier", ctx=ctx)
        _require_str(event, "source_type", ctx=ctx)
        if not isinstance(event.get("event_categories"), list):
            raise ValueError(f"{ctx}: 'event_categories' must be a list")
    return data  # type: ignore[return-value]


def validate_performance_insights(data: dict[str, Any]) -> PerformanceInsightsFixture:
    ctx = "aws_performance_insights.json"
    _require_str(data, "db_instance_identifier", ctx=ctx)
    _require_str(data, "start_time", ctx=ctx)
    _require_str(data, "end_time", ctx=ctx)
    db_load = data.get("db_load")
    if not isinstance(db_load, dict):
        raise ValueError(f"{ctx}: 'db_load' must be an object")
    if not isinstance(db_load.get("timestamps"), list):
        raise ValueError(f"{ctx}: 'db_load.timestamps' must be a list")
    if not isinstance(db_load.get("values"), list):
        raise ValueError(f"{ctx}: 'db_load.values' must be a list")
    if len(db_load["timestamps"]) != len(db_load["values"]):
        raise ValueError(
            f"{ctx}: 'db_load.timestamps' and 'db_load.values' must have the same length"
        )
    if not isinstance(data.get("top_sql"), list):
        raise ValueError(f"{ctx}: 'top_sql' must be a list")
    if not isinstance(data.get("top_wait_events"), list):
        raise ValueError(f"{ctx}: 'top_wait_events' must be a list")
    if not isinstance(data.get("top_users"), list):
        raise ValueError(f"{ctx}: 'top_users' must be a list")
    if not isinstance(data.get("top_hosts"), list):
        raise ValueError(f"{ctx}: 'top_hosts' must be a list")
    return data  # type: ignore[return-value]


def validate_ec2_instances_by_tag(data: dict[str, Any]) -> EC2InstancesByTagFixture:
    ctx = "ec2_instances_by_tag.json"
    instances = data.get("instances")
    if not isinstance(instances, list):
        raise ValueError(f"{ctx}: 'instances' must be a list")
    for i, inst in enumerate(instances):
        ictx = f"{ctx}:instances[{i}]"
        if not isinstance(inst, dict):
            raise ValueError(f"{ictx}: must be an object")
        _require_str(inst, "instance_id", ctx=ictx)
    return data  # type: ignore[return-value]


def validate_elb_target_health(data: dict[str, Any]) -> ELBTargetHealthFixture:
    ctx = "elb_target_health.json"
    if not isinstance(data.get("target_groups"), list):
        raise ValueError(f"{ctx}: 'target_groups' must be a list")
    if not isinstance(data.get("targets"), list):
        raise ValueError(f"{ctx}: 'targets' must be a list")
    for i, target in enumerate(data["targets"]):
        tctx = f"{ctx}:targets[{i}]"
        if not isinstance(target, dict):
            raise ValueError(f"{tctx}: must be an object")
        _require_str(target, "instance_id", ctx=tctx)
        _require_str(target, "state", ctx=tctx)
    return data  # type: ignore[return-value]


def validate_generic_evidence(data: dict[str, Any], *, filename: str) -> GenericEvidenceFixture:
    if not isinstance(data, dict):
        raise ValueError(f"{filename}: expected an object")
    return data  # type: ignore[return-value]


def validate_answer_key(data: dict[str, Any]) -> AnswerKeySchema:
    _require_str(data, "root_cause_category", ctx="answer.yml")
    _require_non_empty_str_list(data, "required_keywords", "answer.yml", required=True)
    _require_str(data, "model_response", ctx="answer.yml")
    for opt_list_field in (
        "forbidden_categories",
        "forbidden_keywords",
        "required_evidence_sources",
        "equivalent_root_cause_categories",
    ):
        val = data.get(opt_list_field)
        if val is not None and not isinstance(val, list):
            raise ValueError(f"answer.yml: '{opt_list_field}' must be a list when present")
    required_sources = data.get("required_evidence_sources")
    if required_sources is not None and required_sources:
        if not all(isinstance(item, str) and item.strip() for item in required_sources):
            raise ValueError(
                "answer.yml: 'required_evidence_sources' must contain only non-empty strings"
            )
        unknown_sources = [item for item in required_sources if item not in VALID_EVIDENCE_SOURCES]
        if unknown_sources:
            raise ValueError(
                "answer.yml: unknown source(s) in required_evidence_sources "
                f"{unknown_sources}; expected subset of {sorted(VALID_EVIDENCE_SOURCES)}"
            )
    equiv = data.get("equivalent_root_cause_categories")
    if (
        equiv is not None
        and equiv
        and not all(isinstance(item, str) and item.strip() for item in equiv)
    ):
        raise ValueError(
            "answer.yml: 'equivalent_root_cause_categories' "
            "must contain only non-empty strings when present"
        )
    trajectory = data.get("optimal_trajectory")
    if trajectory is not None:
        if not isinstance(trajectory, list) or not trajectory:
            raise ValueError(
                "answer.yml: 'optimal_trajectory' must be a non-empty list when present"
            )
        unknown_actions = [a for a in trajectory if a not in VALID_TRAJECTORY_ACTIONS]
        if unknown_actions:
            raise ValueError(
                f"answer.yml: unknown action(s) in optimal_trajectory {unknown_actions}; "
                f"expected subset of {sorted(VALID_TRAJECTORY_ACTIONS)}"
            )
    max_loops = data.get("max_investigation_loops")
    if max_loops is not None and (not isinstance(max_loops, int) or max_loops < 1):
        raise ValueError(
            "answer.yml: 'max_investigation_loops' must be a positive integer when present"
        )
    for axis2_list_field in ("ruling_out_keywords", "required_queries"):
        _require_non_empty_str_list(data, axis2_list_field, "answer.yml")
    golden = data.get("golden_trajectory")
    if golden is not None:
        if not isinstance(golden, dict):
            raise ValueError("answer.yml: 'golden_trajectory' must be an object when present")
        ordered_actions = golden.get("ordered_actions")
        if ordered_actions is not None:
            if (
                not isinstance(ordered_actions, list)
                or not ordered_actions
                or not all(isinstance(action, str) and action.strip() for action in ordered_actions)
            ):
                raise ValueError(
                    "answer.yml: 'golden_trajectory.ordered_actions' must be a non-empty list "
                    "of strings when present"
                )
            unknown_actions = [a for a in ordered_actions if a not in VALID_TRAJECTORY_ACTIONS]
            if unknown_actions:
                raise ValueError(
                    "answer.yml: unknown action(s) in golden_trajectory.ordered_actions "
                    f"{unknown_actions}; expected subset of {sorted(VALID_TRAJECTORY_ACTIONS)}"
                )
        matching = golden.get("matching")
        if matching is not None and matching not in {"strict", "lcs", "set"}:
            raise ValueError(
                "answer.yml: 'golden_trajectory.matching' must be one of "
                "'strict', 'lcs', or 'set' when present"
            )
        for int_field in ("max_edit_distance", "max_extra_actions", "max_redundancy", "max_loops"):
            value = golden.get(int_field)
            if value is not None and (not isinstance(value, int) or value < 0):
                raise ValueError(
                    f"answer.yml: 'golden_trajectory.{int_field}' must be a non-negative integer "
                    "when present"
                )
    return data  # type: ignore[return-value]


def validate_scenario_metadata(data: dict[str, Any]) -> ScenarioMetadataSchema:
    ctx = "scenario.yml"
    for field in (
        "schema_version",
        "scenario_id",
        "engine",
        "engine_version",
        "instance_class",
        "region",
        "db_instance_identifier",
        "failure_mode",
        "severity",
    ):
        _require_str(data, field, ctx=ctx)

    engine = data["engine"]
    if engine not in VALID_ENGINES:
        raise ValueError(
            f"{ctx}: unknown engine {engine!r}; expected one of {sorted(VALID_ENGINES)}"
        )

    failure_mode = data["failure_mode"]
    if failure_mode not in VALID_FAILURE_MODES:
        raise ValueError(
            f"{ctx}: unknown failure_mode {failure_mode!r}; expected one of {sorted(VALID_FAILURE_MODES)}"
        )

    sources = data.get("available_evidence")
    if not isinstance(sources, list) or not sources:
        raise ValueError(f"{ctx}: 'available_evidence' must be a non-empty list")
    unknown = [s for s in sources if s not in VALID_EVIDENCE_SOURCES]
    if unknown:
        raise ValueError(
            f"{ctx}: unknown evidence source(s) {unknown}; expected subset of {sorted(VALID_EVIDENCE_SOURCES)}"
        )

    return data  # type: ignore[return-value]


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _require_str(obj: dict[str, Any], key: str, ctx: str = "") -> None:
    value = obj.get(key)
    prefix = f"{ctx}: " if ctx else ""
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{prefix}missing or empty required string field '{key}'")


def _require_non_empty_str_list(
    obj: dict[str, Any],
    key: str,
    ctx: str,
    *,
    required: bool = False,
) -> None:
    value = obj.get(key)

    if value is None:
        if required:
            raise ValueError(f"{ctx}: '{key}' must be a non-empty list")
        return

    if not isinstance(value, list) or not value:
        raise ValueError(f"{ctx}: '{key}' must be a non-empty list")

    if not all(isinstance(item, str) and item.strip() for item in value):
        raise ValueError(f"{ctx}: all '{key}' entries must be non-empty strings")