#!/usr/bin/env python3
"""
Validate process_spec.json against configuration.json.

Checks:
1. Model type is supported
2. dt > 0
3. level_min < level_max
4. level_init in [level_min, level_max]
5. Signal keys exist in HIL physical_values
6. (Optional) Tick test: run 100 simulation steps and verify bounds

Usage:
    python3 -m tools.validate_process_spec \
        --spec outputs/process_spec.json \
        --config outputs/configuration.json
"""

from __future__ import annotations

import argparse
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Set

from models.process_spec import ProcessSpec


SUPPORTED_MODELS = {"water_tank_v1"}


@dataclass
class ValidationIssue:
    kind: str
    message: str


def extract_hil_physical_value_keys(config: Dict[str, Any]) -> Dict[str, Set[str]]:
    """
    Extract physical_values keys per HIL from configuration.

    Returns: {hil_name: {key1, key2, ...}}
    """
    result: Dict[str, Set[str]] = {}
    for hil in config.get("hils", []):
        name = hil.get("name", "")
        keys: Set[str] = set()
        for pv in hil.get("physical_values", []):
            k = pv.get("name")
            if k:
                keys.add(k)
        result[name] = keys
    return result


def get_all_hil_keys(config: Dict[str, Any]) -> Set[str]:
    """Get union of all HIL physical_values keys."""
    all_keys: Set[str] = set()
    for hil in config.get("hils", []):
        for pv in hil.get("physical_values", []):
            k = pv.get("name")
            if k:
                all_keys.add(k)
    return all_keys


def validate_process_spec(
    spec: ProcessSpec,
    config: Dict[str, Any],
) -> List[ValidationIssue]:
    """Validate ProcessSpec against configuration."""
    issues: List[ValidationIssue] = []

    # 1. Model type supported
    if spec.model not in SUPPORTED_MODELS:
        issues.append(ValidationIssue(
            kind="MODEL",
            message=f"Unsupported model '{spec.model}'. Supported: {SUPPORTED_MODELS}",
        ))

    # 2. dt > 0 (already enforced by Pydantic, but double-check)
    if spec.dt <= 0:
        issues.append(ValidationIssue(
            kind="PARAMS",
            message=f"dt must be > 0, got {spec.dt}",
        ))

    # 3. level_min < level_max
    p = spec.params
    if p.level_min >= p.level_max:
        issues.append(ValidationIssue(
            kind="PARAMS",
            message=f"level_min ({p.level_min}) must be < level_max ({p.level_max})",
        ))

    # 4. level_init in bounds
    if not (p.level_min <= p.level_init <= p.level_max):
        issues.append(ValidationIssue(
            kind="PARAMS",
            message=f"level_init ({p.level_init}) must be in [{p.level_min}, {p.level_max}]",
        ))

    # 5. Signal keys exist in HIL physical_values
    all_hil_keys = get_all_hil_keys(config)
    s = spec.signals

    if s.tank_level_key not in all_hil_keys:
        issues.append(ValidationIssue(
            kind="SIGNALS",
            message=f"tank_level_key '{s.tank_level_key}' not in HIL physical_values. Available: {sorted(all_hil_keys)}",
        ))

    if s.valve_open_key not in all_hil_keys:
        issues.append(ValidationIssue(
            kind="SIGNALS",
            message=f"valve_open_key '{s.valve_open_key}' not in HIL physical_values. Available: {sorted(all_hil_keys)}",
        ))

    if s.level_measured_key not in all_hil_keys:
        issues.append(ValidationIssue(
            kind="SIGNALS",
            message=f"level_measured_key '{s.level_measured_key}' not in HIL physical_values. Available: {sorted(all_hil_keys)}",
        ))

    return issues


def run_tick_test(spec: ProcessSpec, steps: int = 100) -> List[ValidationIssue]:
    """
    Run a pure-Python tick test to verify physics stays bounded.

    Simulates the water tank for `steps` iterations and checks:
    - Level stays in [level_min, level_max]
    - No NaN or Inf values
    """
    issues: List[ValidationIssue] = []

    if spec.model != "water_tank_v1":
        issues.append(ValidationIssue(
            kind="TICK_TEST",
            message=f"Tick test not implemented for model '{spec.model}'",
        ))
        return issues

    p = spec.params
    dt = spec.dt

    # Simulate with valve open
    level = p.level_init
    for i in range(steps):
        q_in = p.q_in_max  # valve open
        q_out = p.k_out * math.sqrt(max(level, 0.0))
        d_level = (q_in - q_out) / p.area * dt
        level = level + d_level

        # Clamp (as the generated code does)
        level = max(p.level_min, min(p.level_max, level))

        # Check for NaN/Inf
        if math.isnan(level) or math.isinf(level):
            issues.append(ValidationIssue(
                kind="TICK_TEST",
                message=f"Level became NaN/Inf at step {i} (valve open)",
            ))
            return issues

    # Check final level is in bounds
    if not (p.level_min <= level <= p.level_max):
        issues.append(ValidationIssue(
            kind="TICK_TEST",
            message=f"Level {level} out of bounds after {steps} steps (valve open)",
        ))

    # Simulate with valve closed (drain only)
    level = p.level_init
    for i in range(steps):
        q_in = 0.0  # valve closed
        q_out = p.k_out * math.sqrt(max(level, 0.0))
        d_level = (q_in - q_out) / p.area * dt
        level = level + d_level
        level = max(p.level_min, min(p.level_max, level))

        if math.isnan(level) or math.isinf(level):
            issues.append(ValidationIssue(
                kind="TICK_TEST",
                message=f"Level became NaN/Inf at step {i} (valve closed)",
            ))
            return issues

    if not (p.level_min <= level <= p.level_max):
        issues.append(ValidationIssue(
            kind="TICK_TEST",
            message=f"Level {level} out of bounds after {steps} steps (valve closed)",
        ))

    return issues


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Validate process_spec.json against configuration.json"
    )
    parser.add_argument(
        "--spec",
        required=True,
        help="Path to process_spec.json",
    )
    parser.add_argument(
        "--config",
        required=True,
        help="Path to configuration.json",
    )
    parser.add_argument(
        "--tick-test",
        action="store_true",
        default=True,
        help="Run tick test (100 steps) to verify physics bounds (default: True)",
    )
    parser.add_argument(
        "--no-tick-test",
        action="store_false",
        dest="tick_test",
        help="Skip tick test",
    )
    args = parser.parse_args()

    spec_path = Path(args.spec)
    config_path = Path(args.config)

    if not spec_path.exists():
        raise SystemExit(f"Spec file not found: {spec_path}")
    if not config_path.exists():
        raise SystemExit(f"Config file not found: {config_path}")

    spec_dict = json.loads(spec_path.read_text(encoding="utf-8"))
    config = json.loads(config_path.read_text(encoding="utf-8"))

    try:
        spec = ProcessSpec.model_validate(spec_dict)
    except Exception as e:
        raise SystemExit(f"Invalid ProcessSpec: {e}")

    print(f"Validating: {spec_path}")
    print(f"Against config: {config_path}")
    print(f"Model: {spec.model}")
    print()

    issues = validate_process_spec(spec, config)

    if args.tick_test:
        print("Running tick test (100 steps)...")
        tick_issues = run_tick_test(spec, steps=100)
        issues.extend(tick_issues)
        if not tick_issues:
            print("  Tick test: PASSED")

    print()
    if issues:
        print(f"VALIDATION FAILED: {len(issues)} issue(s)")
        for issue in issues:
            print(f"  [{issue.kind}] {issue.message}")
        raise SystemExit(1)
    else:
        print("VALIDATION PASSED: process_spec.json is valid")


if __name__ == "__main__":
    main()