tracer-cloud--opensre/tests/synthetic/rds_postgres/scenario_loader.py

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, cast

import yaml

from tests.synthetic.schemas import (
    GoldenTrajectorySchema,
    ScenarioEvidence,
    ScenarioMetadataSchema,
    validate_alert,
    validate_answer_key,
    validate_cloudwatch_metrics,
    validate_ec2_instances_by_tag,
    validate_elb_target_health,
    validate_generic_evidence,
    validate_performance_insights,
    validate_rds_events,
    validate_scenario_metadata,
)

SUITE_DIR = Path(__file__).resolve().parent


@dataclass(frozen=True)
class ScenarioMetadata:
    schema_version: str
    scenario_id: str
    engine: str
    engine_version: str
    instance_class: str
    region: str
    db_instance_identifier: str
    db_cluster: str
    failure_mode: str
    severity: str
    available_evidence: list[str]
    scenario_difficulty: int = 1
    adversarial_signals: list[str] = ()  # type: ignore[assignment]
    depends_on: str = ""


TrajectoryMatching = Literal["strict", "lcs", "set"]


@dataclass(frozen=True)
class GoldenTrajectoryConfig:
    ordered_actions: list[str]
    matching: TrajectoryMatching = "lcs"
    max_edit_distance: int | None = None
    max_extra_actions: int | None = None
    max_redundancy: int | None = None
    max_loops: int | None = None


@dataclass(frozen=True)
class ScenarioAnswerKey:
    root_cause_category: str
    required_keywords: list[str]
    model_response: str
    equivalent_root_cause_categories: tuple[str, ...] = ()
    forbidden_categories: list[str] = ()  # type: ignore[assignment]
    forbidden_keywords: list[str] = ()  # type: ignore[assignment]
    required_evidence_sources: list[str] = ()  # type: ignore[assignment]
    optimal_trajectory: list[str] = ()  # type: ignore[assignment]
    max_investigation_loops: int = 1
    ruling_out_keywords: list[str] = ()  # type: ignore[assignment]
    required_queries: list[str] = ()  # type: ignore[assignment]
    golden_trajectory: GoldenTrajectoryConfig | None = None


@dataclass(frozen=True)
class ScenarioFixture:
    scenario_id: str
    scenario_dir: Path
    alert: dict[str, Any]
    evidence: ScenarioEvidence
    metadata: ScenarioMetadata
    answer_key: ScenarioAnswerKey
    problem_md: str


def _parse_trajectory_matching(value: Any) -> TrajectoryMatching:
    if value in {"strict", "lcs", "set"}:
        return cast(TrajectoryMatching, value)
    raise ValueError(
        "answer.yml: 'golden_trajectory.matching' must be one of "
        "'strict', 'lcs', or 'set' when present"
    )


def _parse_non_negative_int(
    golden_trajectory: GoldenTrajectorySchema | dict[str, Any], field: str
) -> int | None:
    value = golden_trajectory.get(field)
    if value is None:
        return None
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise ValueError(
            f"answer.yml: 'golden_trajectory.{field}' must be a non-negative integer when present"
        )
    return cast(int, value)


def _read_json(path: Path) -> dict[str, Any]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"Expected JSON object in {path}")
    return payload


def _read_yaml(path: Path) -> dict[str, Any]:
    payload = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"Expected YAML object in {path}")
    return payload


# ---------------------------------------------------------------------------
# Base-inheritance helpers
# ---------------------------------------------------------------------------


def _resolve_base_dir(suite_dir: Path, base_id: str) -> Path:
    """Find the base scenario directory by its directory name (e.g. '000-healthy')."""
    base_dir = suite_dir / base_id
    if not base_dir.is_dir():
        raise ValueError(f"Base scenario '{base_id}' not found at {base_dir}")
    base_raw = _read_yaml(base_dir / "scenario.yml")
    if "base" in base_raw:
        raise ValueError(
            f"Chained inheritance is not supported: base scenario '{base_id}' "
            f"itself declares base '{base_raw['base']}'"
        )
    return base_dir


def _merge_scenario_yaml(
    base_raw: dict[str, Any],
    scenario_raw: dict[str, Any],
) -> dict[str, Any]:
    """Shallow-merge scenario overrides on top of base metadata.

    scenario_raw values win. The ``base`` directive is consumed and removed.
    """
    merged = {**base_raw, **{k: v for k, v in scenario_raw.items() if k != "base"}}
    merged.pop("base", None)
    return merged


def _resolve_evidence_path(
    scenario_dir: Path,
    base_dir: Path | None,
    filename: str,
) -> Path:
    """Return the scenario's own evidence file if it exists, otherwise the base's."""
    for search_dir in (scenario_dir, base_dir):
        if search_dir is None:
            continue
        candidate = search_dir / filename
        if candidate.exists():
            return candidate
    raise FileNotFoundError(
        f"Evidence '{filename}' not found in {scenario_dir}"
        + (f" or base {base_dir}" if base_dir else "")
    )


def _has_split_cloudwatch_metrics(scenario_dir: Path) -> bool:
    """Check whether a directory uses per-metric prefixed files."""
    return (scenario_dir / "aws_cloudwatch_metrics_envelope.json").exists()


def _load_cloudwatch_metrics_split(scenario_dir: Path) -> dict[str, Any]:
    """Assemble CloudWatch metrics from prefixed per-metric files.

    Expects ``aws_cloudwatch_metrics_envelope.json`` (shared metadata) and
    ``aws_cloudwatch_metrics_<MetricName>.json`` files in *scenario_dir*.
    """
    envelope = _read_json(scenario_dir / "aws_cloudwatch_metrics_envelope.json")
    prefix = "aws_cloudwatch_metrics_"
    metrics = []
    for f in sorted(scenario_dir.glob(f"{prefix}*.json")):
        if f.name == f"{prefix}envelope.json":
            continue
        metrics.append(_read_json(f))
    envelope["metric_data_results"] = metrics
    return envelope


def _load_cloudwatch_metrics(
    scenario_dir: Path,
    base_dir: Path | None,
) -> dict[str, Any]:
    """Load CloudWatch metrics — consolidated file or per-metric split."""
    # 1. Scenario has a consolidated file
    single = scenario_dir / "aws_cloudwatch_metrics.json"
    if single.is_file():
        return _read_json(single)
    # 2. Scenario has per-metric split files
    if _has_split_cloudwatch_metrics(scenario_dir):
        return _load_cloudwatch_metrics_split(scenario_dir)
    # 3. Fall back to base
    if base_dir is not None:
        base_single = base_dir / "aws_cloudwatch_metrics.json"
        if base_single.is_file():
            return _read_json(base_single)
        if _has_split_cloudwatch_metrics(base_dir):
            return _load_cloudwatch_metrics_split(base_dir)
    raise FileNotFoundError(
        f"CloudWatch metrics not found in {scenario_dir}"
        + (f" or base {base_dir}" if base_dir else "")
    )


# ---------------------------------------------------------------------------
# Parsing helpers
# ---------------------------------------------------------------------------


def _validated_metadata(raw: dict[str, Any]) -> ScenarioMetadata:
    """Validate a (possibly merged) raw dict and return a ScenarioMetadata."""
    validated: ScenarioMetadataSchema = validate_scenario_metadata(raw)
    return ScenarioMetadata(
        schema_version=validated["schema_version"],
        scenario_id=validated["scenario_id"],
        engine=validated["engine"],
        engine_version=validated["engine_version"],
        instance_class=validated["instance_class"],
        region=validated["region"],
        db_instance_identifier=validated["db_instance_identifier"],
        db_cluster=validated.get("db_cluster", ""),
        failure_mode=validated["failure_mode"],
        severity=validated["severity"],
        available_evidence=list(validated["available_evidence"]),
        scenario_difficulty=validated.get("scenario_difficulty", 1),  # type: ignore[arg-type]
        adversarial_signals=list(validated.get("adversarial_signals") or []),
        depends_on=validated.get("depends_on", ""),  # type: ignore[arg-type]
    )


def _parse_scenario_yaml(path: Path) -> tuple[ScenarioMetadata, Path | None]:
    """Parse scenario.yml, resolving base inheritance if declared.

    Returns (metadata, base_dir) where base_dir is the resolved base scenario
    directory, or None if no ``base`` field was declared.
    """
    raw = _read_yaml(path)
    base_id = raw.get("base")
    base_dir: Path | None = None

    if base_id:
        suite_dir = path.parent.parent
        base_dir = _resolve_base_dir(suite_dir, base_id)
        base_raw = _read_yaml(base_dir / "scenario.yml")
        raw = _merge_scenario_yaml(base_raw, raw)

    return _validated_metadata(raw), base_dir


def _parse_answer_yaml(path: Path) -> ScenarioAnswerKey:
    payload = _read_yaml(path)
    validated = validate_answer_key(payload)
    golden_trajectory_raw = validated.get("golden_trajectory")
    golden_trajectory: GoldenTrajectoryConfig | None = None
    if isinstance(golden_trajectory_raw, dict):
        ordered_actions_raw = golden_trajectory_raw.get("ordered_actions")
        if not isinstance(ordered_actions_raw, list) or not ordered_actions_raw:
            raise ValueError(
                "answer.yml: 'golden_trajectory.ordered_actions' must be a non-empty list "
                "of strings when present"
            )
        ordered_actions = [action.strip() for action in ordered_actions_raw]
        matching = _parse_trajectory_matching(golden_trajectory_raw.get("matching", "lcs"))
        golden_trajectory = GoldenTrajectoryConfig(
            ordered_actions=ordered_actions,
            matching=matching,
            max_edit_distance=_parse_non_negative_int(golden_trajectory_raw, "max_edit_distance"),
            max_extra_actions=_parse_non_negative_int(golden_trajectory_raw, "max_extra_actions"),
            max_redundancy=_parse_non_negative_int(golden_trajectory_raw, "max_redundancy"),
            max_loops=_parse_non_negative_int(golden_trajectory_raw, "max_loops"),
        )

    equivalent_raw = validated.get("equivalent_root_cause_categories") or []
    equivalent_root_cause_categories = tuple(
        str(item).strip() for item in equivalent_raw if str(item).strip()
    )

    return ScenarioAnswerKey(
        root_cause_category=validated["root_cause_category"].strip(),
        required_keywords=[k.strip() for k in validated["required_keywords"]],
        model_response=validated["model_response"].strip(),
        equivalent_root_cause_categories=equivalent_root_cause_categories,
        forbidden_categories=list(validated.get("forbidden_categories") or []),
        forbidden_keywords=list(validated.get("forbidden_keywords") or []),
        required_evidence_sources=list(validated.get("required_evidence_sources") or []),
        optimal_trajectory=list(validated.get("optimal_trajectory") or []),
        max_investigation_loops=int(validated.get("max_investigation_loops") or 1),
        ruling_out_keywords=list(validated.get("ruling_out_keywords") or []),
        required_queries=list(validated.get("required_queries") or []),
        golden_trajectory=golden_trajectory,
    )


def _build_problem_md(alert: dict[str, Any], metadata: ScenarioMetadata) -> str:
    title = str(alert.get("title") or metadata.scenario_id)
    annotations = alert.get("commonAnnotations", {}) or {}

    parts = [
        f"# {title}",
        (
            f"Service: RDS {metadata.engine.upper()}"
            f" | Severity: {metadata.severity}"
            f" | Scenario: {metadata.failure_mode}"
        ),
        f"Scenario ID: {metadata.scenario_id}",
        f"DB instance: {metadata.db_instance_identifier}",
    ]

    if metadata.db_cluster:
        parts.append(f"DB cluster: {metadata.db_cluster}")

    summary = annotations.get("summary")
    if summary:
        parts.append(f"\nSummary: {summary}")

    error = annotations.get("error")
    if error and error != summary:
        parts.append(f"\nError: {error}")

    suspected = annotations.get("suspected_symptom")
    if suspected:
        parts.append(f"\nObserved symptom: {suspected}")

    return "\n".join(parts)


def _build_evidence(
    scenario_dir: Path,
    available_evidence: list[str],
    base_dir: Path | None = None,
) -> ScenarioEvidence:
    """Load only the evidence sources declared in scenario.yml:available_evidence.

    When *base_dir* is set, evidence files missing from *scenario_dir* are
    resolved from the base scenario directory (file-level fallback).
    """
    aws_cloudwatch_metrics = None
    aws_rds_events = None
    aws_performance_insights = None
    ec2_instances_by_tag = None
    elb_target_health = None
    k8s_events = None
    k8s_pod_metrics = None
    k8s_node_metrics = None
    k8s_dns_metrics = None
    k8s_mesh_metrics = None
    k8s_rollout = None

    if "aws_cloudwatch_metrics" in available_evidence:
        raw = _load_cloudwatch_metrics(scenario_dir, base_dir)
        aws_cloudwatch_metrics = validate_cloudwatch_metrics(raw)

    if "aws_rds_events" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "aws_rds_events.json")
        raw_events = validate_rds_events(_read_json(path))
        aws_rds_events = raw_events.get("events", [])

    if "aws_performance_insights" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "aws_performance_insights.json")
        aws_performance_insights = validate_performance_insights(_read_json(path))

    if "ec2_instances_by_tag" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "ec2_instances_by_tag.json")
        ec2_instances_by_tag = validate_ec2_instances_by_tag(_read_json(path))

    if "elb_target_health" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "elb_target_health.json")
        elb_target_health = validate_elb_target_health(_read_json(path))

    if "k8s_events" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_events.json")
        k8s_events = validate_generic_evidence(_read_json(path), filename="k8s_events.json")

    if "k8s_pod_metrics" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_pod_metrics.json")
        k8s_pod_metrics = validate_generic_evidence(
            _read_json(path), filename="k8s_pod_metrics.json"
        )

    if "k8s_node_metrics" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_node_metrics.json")
        k8s_node_metrics = validate_generic_evidence(
            _read_json(path), filename="k8s_node_metrics.json"
        )

    if "k8s_dns_metrics" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_dns_metrics.json")
        k8s_dns_metrics = validate_generic_evidence(
            _read_json(path), filename="k8s_dns_metrics.json"
        )

    if "k8s_mesh_metrics" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_mesh_metrics.json")
        k8s_mesh_metrics = validate_generic_evidence(
            _read_json(path), filename="k8s_mesh_metrics.json"
        )

    if "k8s_rollout" in available_evidence:
        path = _resolve_evidence_path(scenario_dir, base_dir, "k8s_rollout.json")
        k8s_rollout = validate_generic_evidence(_read_json(path), filename="k8s_rollout.json")

    return ScenarioEvidence(
        aws_cloudwatch_metrics=aws_cloudwatch_metrics,
        aws_rds_events=aws_rds_events,
        aws_performance_insights=aws_performance_insights,
        ec2_instances_by_tag=ec2_instances_by_tag,
        elb_target_health=elb_target_health,
        k8s_events=k8s_events,
        k8s_pod_metrics=k8s_pod_metrics,
        k8s_node_metrics=k8s_node_metrics,
        k8s_dns_metrics=k8s_dns_metrics,
        k8s_mesh_metrics=k8s_mesh_metrics,
        k8s_rollout=k8s_rollout,
    )


def _is_schema_v3(schema_version: str) -> bool:
    normalized = schema_version.strip().lower().replace("-", "_")
    return normalized in {"schema_v3", "v3", "3", "3.0"}


def _is_complex_scenario(metadata: ScenarioMetadata) -> bool:
    return metadata.scenario_difficulty >= 3


def _validate_schema_specific_answer_requirements(
    metadata: ScenarioMetadata,
    answer_key: ScenarioAnswerKey,
) -> None:
    if (
        _is_schema_v3(metadata.schema_version)
        and _is_complex_scenario(metadata)
        and not answer_key.required_evidence_sources
    ):
        raise ValueError(
            "answer.yml: 'required_evidence_sources' must be a non-empty list "
            "for schema_v3 complex scenarios (scenario_difficulty >= 3)"
        )


def load_scenario(scenario_dir: Path) -> ScenarioFixture:
    metadata, base_dir = _parse_scenario_yaml(scenario_dir / "scenario.yml")

    alert_path = _resolve_evidence_path(scenario_dir, base_dir, "alert.json")
    alert = cast(dict[str, Any], validate_alert(_read_json(alert_path)))

    evidence = _build_evidence(scenario_dir, metadata.available_evidence, base_dir)
    answer_key = _parse_answer_yaml(scenario_dir / "answer.yml")
    _validate_schema_specific_answer_requirements(metadata, answer_key)
    problem_md = _build_problem_md(alert, metadata)

    return ScenarioFixture(
        scenario_id=scenario_dir.name,
        scenario_dir=scenario_dir,
        alert=alert,
        evidence=evidence,
        metadata=metadata,
        answer_key=answer_key,
        problem_md=problem_md,
    )


def load_all_scenarios(root_dir: Path | None = None) -> list[ScenarioFixture]:
    base_dir = root_dir or SUITE_DIR
    scenario_dirs = sorted(
        path for path in base_dir.iterdir() if path.is_dir() and path.name[:3].isdigit()
    )
    return [load_scenario(path) for path in scenario_dirs]