"""End-to-End Simulation Runner.

Executes the full pipeline from CNL spec to Nengo simulation
and produces a validation report.

Usage:
    python run_simulation.py reflex_arc.cnl
"""

import json
import os
import re
import subprocess
import sys
import tempfile
import time
from pathlib import Path

# Ensure all neurocnl subdirectories are on sys.path when this script is run
# directly (e.g. `python neurocnl/simulation/run_simulation.py ...`).
_neurocnl_root = Path(__file__).resolve().parent.parent
for _subdir in ["", "cnl", "layers", "generation", "simulation"]:
    _p = str(_neurocnl_root / _subdir) if _subdir else str(_neurocnl_root)
    if _p not in sys.path:
        sys.path.insert(0, _p)

import nengo
import numpy as np

from assertion_generator import generate_assertions
from cnl_parser import ParseError, parse
from layer1_validator import validate
from nengo_generator import GeneratorError, generate
from utils import extract_numeric


def load_spec(spec_path: str) -> list[str]:
    """Load CNL sentences from a spec file (one sentence per line)."""
    with open(spec_path) as f:
        lines = [line.strip() for line in f if line.strip()]
    return lines


def run_pipeline(spec_path: str) -> dict:
    """Execute the full simulation pipeline.

    Parameters
    ----------
    spec_path : str
        Path to a .cnl spec file.

    Returns
    -------
    dict
        Simulation report.
    """
    report = {
        "spec_file": spec_path,
        "layer1_validation": None,
        "simulation_duration": 0.0,
        "mujoco_steps": 0,
        "assertions_passed": 0,
        "assertions_failed": 0,
        "overall_pass": False,
    }

    # Step 1-2: Parse the spec
    sentences = load_spec(spec_path)
    try:
        parsed_specs = [parse(s) for s in sentences]
    except ParseError as e:
        report["error"] = f"Parse error: {e}"
        return report

    # Step 3: Validate against Layer 1
    neuron_params = {
        "threshold": 1.0,
        "resting_potential": 0.0,
        "refractory_period": 0.002,
        "tau": 0.02,
        "reset_potential": 0.0,
        "current_voltage": 0.5,
    }

    # Extract parameters from parsed specs
    _CONCEPT_TO_PARAM = {
        "refractory_period": "refractory_period",
        "membrane_potential_decay": "tau",
    }
    for spec in parsed_specs:
        param_key = _CONCEPT_TO_PARAM.get(spec["concept"])
        if param_key is not None:
            val = extract_numeric(spec["condition"])
            if val is not None:
                neuron_params[param_key] = val

    validation = validate(parsed_specs, neuron_params)
    report["layer1_validation"] = validation

    if not validation["overall"]:
        report["error"] = "Layer 1 validation failed"
        return report

    # Step 4: Generate Nengo network
    try:
        net = generate(parsed_specs, neuron_params)
    except GeneratorError as e:
        report["error"] = f"Generator error: {e}"
        return report

    # Step 5: Run Nengo simulation for 1 second
    start_time = time.time()
    with nengo.Simulator(net) as sim:
        sim.run(1.0)
    sim_duration = time.time() - start_time
    report["simulation_duration"] = round(sim_duration, 3)

    motor_output = sim.data[net.motor_probe]

    # Step 6-7: MuJoCo integration
    # Use direct MuJoCo API since gymnasium may not be installed
    mujoco_steps = 0
    try:
        import mujoco

        # Create a simple inverted pendulum model
        xml = """
        <mujoco>
          <worldbody>
            <body name="cart" pos="0 0 0">
              <joint name="slide" type="slide" axis="1 0 0"/>
              <geom type="box" size="0.2 0.1 0.05" mass="1"/>
              <body name="pole" pos="0 0 0.6">
                <joint name="hinge" type="hinge" axis="0 1 0"/>
                <geom type="capsule" fromto="0 0 0 0 0 0.6" size="0.02" mass="0.1"/>
              </body>
            </body>
          </worldbody>
          <actuator>
            <motor joint="slide" gear="100"/>
          </actuator>
        </mujoco>
        """
        model = mujoco.MjModel.from_xml_string(xml)
        data = mujoco.MjData(model)

        # Run 200 steps using motor output as control signal
        for step_i in range(200):
            # Map simulation output to action
            t_idx = min(step_i, len(motor_output) - 1)
            action = float(np.clip(motor_output[t_idx, 0], -1.0, 1.0))
            data.ctrl[0] = action
            mujoco.mj_step(model, data)
            mujoco_steps += 1

            # Check termination (pole angle)
            pole_angle = data.qpos[1] if len(data.qpos) > 1 else 0.0
            if abs(pole_angle) > 0.2:
                break

    except Exception as e:
        # MuJoCo is optional; record what we got
        report["mujoco_note"] = f"MuJoCo integration: {e}"

    report["mujoco_steps"] = mujoco_steps

    # Step 8: Generate and run assertions
    with tempfile.TemporaryDirectory() as tmpdir:
        assertion_path = os.path.join(tmpdir, "test_layer3_assertions.py")
        generate_assertions(parsed_specs, output_path=assertion_path)

        pytest_result = subprocess.run(
            [sys.executable, "-m", "pytest", assertion_path, "-v", "--tb=short", "-p", "no:nengo"],
            capture_output=True,
            text=True,
        )

    # Parse pytest output for pass/fail counts
    for line in pytest_result.stdout.splitlines():
        if "passed" in line:
            m = re.search(r"(\d+)\s+passed", line)
            if m:
                report["assertions_passed"] = int(m.group(1))
        if "failed" in line:
            m = re.search(r"(\d+)\s+failed", line)
            if m:
                report["assertions_failed"] = int(m.group(1))

    # Surface pytest errors (non-zero exit = collection error, import error, etc.)
    if pytest_result.returncode not in (0, 1):
        report["assertions_error"] = (
            pytest_result.stderr or pytest_result.stdout
        )[:2000]

    # Determine overall pass
    report["overall_pass"] = (
        validation["overall"]
        and report["assertions_failed"] == 0
        and report["assertions_passed"] > 0
        and pytest_result.returncode in (0, 1)
    )

    return report


def main():
    if len(sys.argv) < 2:
        print("Usage: python run_simulation.py <spec_file.cnl>")
        sys.exit(1)

    spec_path = sys.argv[1]
    report = run_pipeline(spec_path)

    # Write report
    report_path = "simulation_report.json"
    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)

    print(json.dumps(report, indent=2))

    if report["overall_pass"]:
        print(f"\n✓ OVERALL PASS — report written to {report_path}")
        sys.exit(0)
    else:
        print(f"\n✗ OVERALL FAIL — report written to {report_path}")
        sys.exit(1)


if __name__ == "__main__":
    main()
