#!/usr/bin/env python3
"""
Compile control_plan.json into deterministic HIL logic files.

Input: control_plan.json (ControlPlan schema)
Output: Python HIL logic files (one per HIL in the plan)

Usage:
    python3 -m tools.compile_control_plan \
        --control-plan outputs/control_plan.json \
        --out outputs/scenario_run/logic

With validation against config:
    python3 -m tools.compile_control_plan \
        --control-plan outputs/control_plan.json \
        --out outputs/scenario_run/logic \
        --config outputs/configuration.json

Validation only (no file generation):
    python3 -m tools.compile_control_plan \
        --control-plan outputs/control_plan.json \
        --validate-only

The generated HIL logic follows the ICS-SimLab contract:
    def logic(physical_values):
        # Initialize from plan.init
        # Optional warmup sleep
        # while True: run tasks (threaded if >1)
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Union

from models.control_plan import (
    Action,
    AddAction,
    ControlPlan,
    ControlPlanHIL,
    GaussianProfile,
    IfAction,
    LoopTask,
    PlaybackTask,
    RampProfile,
    SetAction,
    StepProfile,
    Task,
)
from tools.safe_eval import (
    UnsafeExpressionError,
    extract_variable_names,
    generate_python_code,
    validate_expression,
)


class CompilationError(Exception):
    """Raised when control plan compilation fails."""
    pass


class ValidationError(Exception):
    """Raised when control plan validation fails."""
    pass


def get_hil_physical_values_keys(config: dict, hil_name: str) -> Tuple[Set[str], Set[str]]:
    """
    Extract (input_keys, output_keys) for a specific HIL from config.

    Returns tuple of (set of input keys, set of output keys).
    """
    input_keys: Set[str] = set()
    output_keys: Set[str] = set()

    for hil in config.get("hils", []):
        if hil.get("name") == hil_name:
            for pv in hil.get("physical_values", []):
                key = pv.get("name")
                io = pv.get("io")
                if key:
                    if io == "input":
                        input_keys.add(key)
                    elif io == "output":
                        output_keys.add(key)
            break

    return input_keys, output_keys


def validate_control_plan(
    plan: ControlPlan,
    config: Optional[dict] = None
) -> List[str]:
    """
    Validate a control plan for errors.

    Checks:
    1. All expressions are syntactically valid and safe
    2. All variables in expressions exist in init/params or config physical_values
    3. All set/add targets are valid output keys (if config provided)

    Returns:
        List of error messages (empty if valid)
    """
    errors: List[str] = []

    for hil in plan.hils:
        hil_name = hil.name

        # Build available namespace: init + params
        available_vars: Set[str] = set(hil.init.keys())
        if hil.params:
            available_vars.update(hil.params.keys())

        # Add config physical_values if provided
        config_input_keys: Set[str] = set()
        config_output_keys: Set[str] = set()
        if config:
            config_input_keys, config_output_keys = get_hil_physical_values_keys(config, hil_name)
            available_vars.update(config_input_keys)
            available_vars.update(config_output_keys)

        # Validate tasks
        for task in hil.tasks:
            if isinstance(task, LoopTask):
                # Validate all actions in the loop
                action_errors = _validate_actions(
                    task.actions,
                    available_vars,
                    config_output_keys if config else None,
                    hil_name,
                    task.name
                )
                errors.extend(action_errors)

            elif isinstance(task, PlaybackTask):
                # Validate target variable
                target = task.target
                if target not in available_vars:
                    errors.append(
                        f"[{hil_name}/{task.name}] Playback target '{target}' not defined in init/params/physical_values"
                    )
                if config and config_output_keys and target not in config_output_keys:
                    errors.append(
                        f"[{hil_name}/{task.name}] Playback target '{target}' is not an output in config"
                    )

    return errors


def _validate_actions(
    actions: List[Action],
    available_vars: Set[str],
    output_keys: Optional[Set[str]],
    hil_name: str,
    task_name: str
) -> List[str]:
    """Validate a list of actions recursively."""
    errors: List[str] = []
    prefix = f"[{hil_name}/{task_name}]"

    for action in actions:
        if isinstance(action, SetAction):
            var, expr = action.set
            # Validate expression
            try:
                validate_expression(expr)
            except (SyntaxError, UnsafeExpressionError) as e:
                errors.append(f"{prefix} Invalid expression in set({var}): {e}")
                continue

            # Check referenced variables exist
            refs = extract_variable_names(expr)
            undefined = refs - available_vars
            if undefined:
                errors.append(f"{prefix} Undefined variables in set({var}): {undefined}")

            # Check target is writable
            if output_keys is not None and var not in output_keys and var not in available_vars:
                errors.append(f"{prefix} set target '{var}' is not defined in init/params/outputs")

        elif isinstance(action, AddAction):
            var, expr = action.add
            # Validate expression
            try:
                validate_expression(expr)
            except (SyntaxError, UnsafeExpressionError) as e:
                errors.append(f"{prefix} Invalid expression in add({var}): {e}")
                continue

            # Check referenced variables exist
            refs = extract_variable_names(expr)
            undefined = refs - available_vars
            if undefined:
                errors.append(f"{prefix} Undefined variables in add({var}): {undefined}")

            # Check target is writable
            if output_keys is not None and var not in output_keys and var not in available_vars:
                errors.append(f"{prefix} add target '{var}' is not defined in init/params/outputs")

        elif isinstance(action, IfAction):
            cond = action.if_
            # Validate condition
            try:
                validate_expression(cond)
            except (SyntaxError, UnsafeExpressionError) as e:
                errors.append(f"{prefix} Invalid condition: {e}")
                continue

            # Check referenced variables
            refs = extract_variable_names(cond)
            undefined = refs - available_vars
            if undefined:
                errors.append(f"{prefix} Undefined variables in condition: {undefined}")

            # Recursively validate then/else actions
            errors.extend(_validate_actions(action.then, available_vars, output_keys, hil_name, task_name))
            if action.else_:
                errors.extend(_validate_actions(action.else_, available_vars, output_keys, hil_name, task_name))

    return errors


def _indent(code: str, level: int = 1) -> str:
    """Indent code by the given number of levels (4 spaces each)."""
    prefix = "    " * level
    return "\n".join(prefix + line if line else line for line in code.split("\n"))


def _compile_action(action: Action, indent_level: int, pv_var: str = "pv") -> str:
    """Compile a single action to Python code."""
    lines: List[str] = []
    indent = "    " * indent_level

    if isinstance(action, SetAction):
        var, expr = action.set
        py_expr = generate_python_code(expr, pv_var)
        lines.append(f"{indent}{pv_var}['{var}'] = {py_expr}")

    elif isinstance(action, AddAction):
        var, expr = action.add
        py_expr = generate_python_code(expr, pv_var)
        lines.append(f"{indent}{pv_var}['{var}'] = {pv_var}.get('{var}', 0) + ({py_expr})")

    elif isinstance(action, IfAction):
        cond = action.if_
        py_cond = generate_python_code(cond, pv_var)
        lines.append(f"{indent}if {py_cond}:")
        for a in action.then:
            lines.append(_compile_action(a, indent_level + 1, pv_var))
        if action.else_:
            lines.append(f"{indent}else:")
            for a in action.else_:
                lines.append(_compile_action(a, indent_level + 1, pv_var))

    return "\n".join(lines)


def _compile_loop_task(task: LoopTask, pv_var: str = "pv") -> str:
    """Compile a loop task to a function definition."""
    lines: List[str] = []
    func_name = f"_task_{task.name.replace('-', '_').replace(' ', '_')}"

    lines.append(f"def {func_name}({pv_var}):")
    lines.append(f'    """Loop task: {task.name} (dt={task.dt_s}s)"""')
    lines.append("    while True:")

    # Compile actions
    for action in task.actions:
        lines.append(_compile_action(action, 2, pv_var))

    lines.append(f"        time.sleep({task.dt_s})")

    return "\n".join(lines)


def _compile_playback_task(task: PlaybackTask, pv_var: str = "pv") -> str:
    """Compile a playback task to a function definition."""
    lines: List[str] = []
    func_name = f"_task_{task.name.replace('-', '_').replace(' ', '_')}"
    profile = task.profile

    lines.append(f"def {func_name}({pv_var}):")
    lines.append(f'    """Playback task: {task.name} (dt={task.dt_s}s)"""')

    # Generate profile data
    if isinstance(profile, GaussianProfile):
        lines.append(f"    # Gaussian profile: height={profile.height}, mean={profile.mean}, std={profile.std}, entries={profile.entries}")
        lines.append(f"    _profile_height = {profile.height}")
        lines.append(f"    _profile_mean = {profile.mean}")
        lines.append(f"    _profile_std = {profile.std}")
        lines.append(f"    _profile_entries = {profile.entries}")
        lines.append("    _profile_idx = 0")
        lines.append("    while True:")
        lines.append("        _value = _profile_height + random.gauss(_profile_mean, _profile_std)")
        lines.append(f"        {pv_var}['{task.target}'] = _value")
        lines.append(f"        time.sleep({task.dt_s})")
        if not task.repeat:
            lines.append("        _profile_idx += 1")
            lines.append("        if _profile_idx >= _profile_entries:")
            lines.append("            break")

    elif isinstance(profile, RampProfile):
        lines.append(f"    # Ramp profile: start={profile.start}, end={profile.end}, entries={profile.entries}")
        lines.append(f"    _profile_start = {profile.start}")
        lines.append(f"    _profile_end = {profile.end}")
        lines.append(f"    _profile_entries = {profile.entries}")
        lines.append("    _profile_idx = 0")
        lines.append("    while True:")
        lines.append("        _t = _profile_idx / max(1, _profile_entries - 1)")
        lines.append("        _value = _profile_start + (_profile_end - _profile_start) * _t")
        lines.append(f"        {pv_var}['{task.target}'] = _value")
        lines.append(f"        time.sleep({task.dt_s})")
        lines.append("        _profile_idx += 1")
        if task.repeat:
            lines.append("        if _profile_idx >= _profile_entries:")
            lines.append("            _profile_idx = 0")
        else:
            lines.append("        if _profile_idx >= _profile_entries:")
            lines.append("            break")

    elif isinstance(profile, StepProfile):
        values_str = repr(profile.values)
        lines.append(f"    # Step profile: values={values_str}")
        lines.append(f"    _profile_values = {values_str}")
        lines.append("    _profile_idx = 0")
        lines.append("    while True:")
        lines.append("        _value = _profile_values[_profile_idx % len(_profile_values)]")
        lines.append(f"        {pv_var}['{task.target}'] = _value")
        lines.append(f"        time.sleep({task.dt_s})")
        lines.append("        _profile_idx += 1")
        if not task.repeat:
            lines.append("        if _profile_idx >= len(_profile_values):")
            lines.append("            break")

    return "\n".join(lines)


def compile_hil(hil: ControlPlanHIL, config_physical_values: Optional[Set[str]] = None) -> str:
    """
    Compile a single HIL control plan to Python code.

    Args:
        hil: The HIL control plan to compile
        config_physical_values: Optional set of physical_value keys from config.
            If provided, ensures ALL keys are initialized (not just plan.init keys).
            Keys in plan.init use their init value; others use 0 as default.

    Returns:
        Python code string for the HIL logic file
    """
    lines: List[str] = []

    # Header
    lines.append('"""')
    lines.append(f"HIL logic for {hil.name}: ControlPlan v0.1 compiled.")
    lines.append("")
    lines.append("Autogenerated by ics-simlab-config-gen (compile_control_plan).")
    lines.append("DO NOT EDIT - regenerate from control_plan.json instead.")
    lines.append('"""')
    lines.append("")

    # Imports
    lines.append("import random")
    lines.append("import time")
    if len(hil.tasks) > 1:
        lines.append("import threading")
    lines.append("")

    # Helper: clamp function
    lines.append("")
    lines.append("def clamp(x, lo, hi):")
    lines.append('    """Clamp x to [lo, hi]."""')
    lines.append("    return lo if x < lo else hi if x > hi else x")
    lines.append("")

    # Compile each task to a function
    for task in hil.tasks:
        lines.append("")
        if isinstance(task, LoopTask):
            lines.append(_compile_loop_task(task, "pv"))
        elif isinstance(task, PlaybackTask):
            lines.append(_compile_playback_task(task, "pv"))
    lines.append("")

    # Main logic function
    lines.append("")
    lines.append("def logic(physical_values):")
    lines.append('    """')
    lines.append(f"    HIL logic entry point for {hil.name}.")
    lines.append("")
    lines.append("    ICS-SimLab calls this once and expects it to run forever.")
    lines.append('    """')

    # === CRITICAL: Initialize physical_values BEFORE any alias or threads ===
    # Use setdefault directly on physical_values (not an alias) so that
    # tools.validate_logic --check-hil-init can detect initialization via AST.
    lines.append("")
    lines.append("    # === Initialize physical values (validator-compatible) ===")

    # Determine all keys that need initialization
    all_keys_to_init: Set[str] = set(hil.init.keys())
    if config_physical_values:
        all_keys_to_init = all_keys_to_init | config_physical_values

    # Emit setdefault for ALL keys, using plan.init value if available, else 0
    for key in sorted(all_keys_to_init):
        if key in hil.init:
            value = hil.init[key]
            if isinstance(value, bool):
                py_val = "True" if value else "False"
            else:
                py_val = repr(value)
        else:
            # Key from config not in plan.init - use 0 as default
            py_val = "0"
        lines.append(f"    physical_values.setdefault('{key}', {py_val})")
    lines.append("")

    # Now create alias for rest of generated code
    lines.append("    pv = physical_values  # Alias for generated code")
    lines.append("")

    # Params as local variables (for documentation; they're used via pv.get)
    if hil.params:
        lines.append("    # === Parameters (read-only constants) ===")
        for key, value in hil.params.items():
            if isinstance(value, bool):
                py_val = "True" if value else "False"
            else:
                py_val = repr(value)
            lines.append(f"    pv['{key}'] = {py_val}  # param")
        lines.append("")

    # Warmup sleep
    if hil.warmup_s:
        lines.append(f"    # === Warmup delay ===")
        lines.append(f"    time.sleep({hil.warmup_s})")
        lines.append("")

    # Start tasks
    if len(hil.tasks) == 1:
        # Single task: just call it directly
        task = hil.tasks[0]
        func_name = f"_task_{task.name.replace('-', '_').replace(' ', '_')}"
        lines.append(f"    # === Run single task ===")
        lines.append(f"    {func_name}(pv)")
    else:
        # Multiple tasks: use threading
        lines.append("    # === Start tasks in threads ===")
        lines.append("    threads = []")
        for task in hil.tasks:
            func_name = f"_task_{task.name.replace('-', '_').replace(' ', '_')}"
            lines.append(f"    t = threading.Thread(target={func_name}, args=(pv,), daemon=True)")
            lines.append("    t.start()")
            lines.append("    threads.append(t)")
        lines.append("")
        lines.append("    # Wait for all threads (they run forever)")
        lines.append("    for t in threads:")
        lines.append("        t.join()")

    lines.append("")

    return "\n".join(lines)


def compile_control_plan(
    plan: ControlPlan,
    config: Optional[dict] = None,
) -> Dict[str, str]:
    """
    Compile a control plan to HIL logic files.

    Args:
        plan: The ControlPlan to compile
        config: Optional configuration.json dict. If provided, ensures ALL
            physical_values declared in config for each HIL are initialized.

    Returns:
        Dict mapping HIL name to Python code string
    """
    result: Dict[str, str] = {}
    for hil in plan.hils:
        # Extract config physical_values for this HIL if config provided
        config_pv: Optional[Set[str]] = None
        if config:
            input_keys, output_keys = get_hil_physical_values_keys(config, hil.name)
            config_pv = input_keys | output_keys
        result[hil.name] = compile_hil(hil, config_physical_values=config_pv)
    return result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compile control_plan.json into HIL logic Python files"
    )
    parser.add_argument(
        "--control-plan",
        required=True,
        help="Path to control_plan.json",
    )
    parser.add_argument(
        "--out",
        default=None,
        help="Output directory for HIL logic .py files",
    )
    parser.add_argument(
        "--config",
        default=None,
        help="Path to configuration.json (for validation)",
    )
    parser.add_argument(
        "--validate-only",
        action="store_true",
        help="Only validate, don't generate files",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite existing output files",
    )
    args = parser.parse_args()

    plan_path = Path(args.control_plan)
    out_dir = Path(args.out) if args.out else None
    config_path = Path(args.config) if args.config else None

    if not plan_path.exists():
        raise SystemExit(f"Control plan not found: {plan_path}")
    if config_path and not config_path.exists():
        raise SystemExit(f"Config file not found: {config_path}")
    if not args.validate_only and not out_dir:
        raise SystemExit("--out is required unless using --validate-only")

    # Load plan
    plan_dict = json.loads(plan_path.read_text(encoding="utf-8"))
    plan = ControlPlan.model_validate(plan_dict)

    print(f"Loaded control plan: {plan_path}")
    print(f"  Version: {plan.version}")
    print(f"  HILs: {[h.name for h in plan.hils]}")

    # Load config for validation
    config: Optional[dict] = None
    if config_path:
        config = json.loads(config_path.read_text(encoding="utf-8"))
        print(f"  Config: {config_path}")

    # Validate
    errors = validate_control_plan(plan, config)
    if errors:
        print(f"\nValidation FAILED ({len(errors)} errors):")
        for err in errors:
            print(f"  - {err}")
        raise SystemExit(1)
    else:
        print("  Validation: OK")

    if args.validate_only:
        print("\nValidation only mode, no files generated.")
        return

    # Compile (pass config to ensure all physical_values are initialized)
    hil_code = compile_control_plan(plan, config=config)

    # Write output files
    assert out_dir is not None
    out_dir.mkdir(parents=True, exist_ok=True)

    for hil_name, code in hil_code.items():
        # Use hil_name as filename (sanitized)
        safe_name = hil_name.replace(" ", "_").replace("-", "_")
        out_file = out_dir / f"{safe_name}.py"

        if out_file.exists() and not args.overwrite:
            raise SystemExit(f"Output file exists: {out_file} (use --overwrite)")

        out_file.write_text(code, encoding="utf-8")
        print(f"Wrote: {out_file}")

    print(f"\nCompiled {len(hil_code)} HIL logic file(s) to {out_dir}")


if __name__ == "__main__":
    main()