import json
import struct
import wave
from pathlib import Path

import pytest

from demos.grownet_trace_main import run_for_test
from runner.experiment_loader import apply_cli_overrides, apply_trace_config_defaults, load_experiment_config_file
from runner.input_loaders import load_input_schedule
from runner.scaffold_builder import build_regions_from_scaffold
from debug.trace_config import TraceConfig


REPO_ROOT = Path(__file__).resolve().parents[3]


def test_experiment_config_parses_checked_in_dual_scalar_example():
    config_path = REPO_ROOT / "experiments" / "dual_scalar_sequence.json"
    experiment_config = load_experiment_config_file(str(config_path))

    assert experiment_config.name == "dual_scalar_sequence"
    assert experiment_config.input_space.rank == 1
    assert list(experiment_config.input_space.shape) == [2]
    assert list(experiment_config.input_space.ports) == ["a", "b"]
    assert len(experiment_config.regions) == 1
    assert len(experiment_config.regions[0].layers) == 2


def test_scaffold_builder_binds_scalar_ports_for_configured_region():
    config_path = REPO_ROOT / "experiments" / "dual_scalar_sequence.json"
    experiment_config = load_experiment_config_file(str(config_path))

    built_regions = build_regions_from_scaffold(experiment_config.regions, experiment_config.input_space)

    assert len(built_regions) == 1
    assert set(built_regions[0].input_edges.keys()) == {"a", "b"}
    assert len(built_regions[0].get_layers()) >= 4


def test_sequence_loader_reads_file_backed_values(tmp_path):
    input_path = tmp_path / "sequence.json"
    input_path.write_text(json.dumps([0.1, 0.2, 0.3]), encoding="utf-8")

    config_path = tmp_path / "sequence_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "file_sequence",
                "runtime": {"ticks": 3},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "path": "sequence.json"},
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))
    loaded_schedule = load_input_schedule(experiment_config)

    assert loaded_schedule.tick_payloads == [{"a": 0.1}, {"a": 0.2}, {"a": 0.3}]


def test_tensor_loader_accepts_inline_2d_tensor_example():
    config_path = REPO_ROOT / "experiments" / "image_tensor.json"
    experiment_config = load_experiment_config_file(str(config_path))

    loaded_schedule = load_input_schedule(experiment_config)

    assert len(loaded_schedule.tick_payloads) == 1
    assert loaded_schedule.tick_payloads[0] == [[0.9, 0.0], [0.0, 0.4]]


def test_tensor_loader_accepts_rank_three_tensor(tmp_path):
    config_path = tmp_path / "tensor_3d.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "tensor_3d_probe",
                "input_space": {"rank": 3, "shape": [2, 1, 2], "ports": [{"name": "volume"}]},
                "network": {
                    "regions": [
                        {
                            "name": "volume_region",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ],
                            "input_bindings": {"volume": [0]},
                        }
                    ]
                },
                "input": {
                    "type": "tensor",
                    "values": [
                        [[0.1, 0.2]],
                        [[0.3, 0.4]]
                    ]
                },
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))
    loaded_schedule = load_input_schedule(experiment_config)

    assert len(loaded_schedule.tick_payloads) == 1
    assert loaded_schedule.tick_payloads[0][1][0][1] == 0.4


def test_image_adapter_loads_small_pgm_and_resizes(tmp_path):
    image_path = tmp_path / "sample.pgm"
    image_path.write_text("P2\n2 2\n255\n0 255\n128 64\n", encoding="utf-8")

    config_path = tmp_path / "image_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "image_probe",
                "input_space": {"rank": 2, "shape": [1, 1], "ports": [{"name": "image"}]},
                "network": {
                    "regions": [
                        {
                            "name": "vision",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ],
                            "input_bindings": {"image": [0]},
                        }
                    ]
                },
                "input": {"type": "image", "path": "sample.pgm", "mode": "grayscale", "resize": [1, 1]},
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))
    loaded_schedule = load_input_schedule(experiment_config)

    assert len(loaded_schedule.tick_payloads) == 1
    assert loaded_schedule.tick_payloads[0] == [[0.0]]


def test_audio_adapter_loads_wav_windows(tmp_path):
    audio_path = tmp_path / "sample.wav"
    with wave.open(str(audio_path), "wb") as wave_handle:
        wave_handle.setnchannels(1)
        wave_handle.setsampwidth(2)
        wave_handle.setframerate(16000)
        samples = [0, 1000, -1000, 2000, -2000, 500, -500, 0]
        wave_handle.writeframes(struct.pack("<8h", *samples))

    config_path = tmp_path / "audio_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "audio_probe",
                "input_space": {"rank": 1, "shape": [4], "ports": [{"name": "audio"}]},
                "network": {
                    "regions": [
                        {
                            "name": "hearing",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {
                    "type": "audio",
                    "path": "sample.wav",
                    "sample_rate": 16000,
                    "frame_size": 4,
                    "hop_size": 2
                },
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))
    loaded_schedule = load_input_schedule(experiment_config)

    assert len(loaded_schedule.tick_payloads) == 3
    assert all(len(frame_values) == 4 for frame_values in loaded_schedule.tick_payloads)


def test_rank_validation_rejects_multi_port_rank_two_config(tmp_path):
    config_path = tmp_path / "invalid_rank.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "invalid_rank",
                "input_space": {
                    "rank": 2,
                    "shape": [2, 2],
                    "ports": [{"name": "left"}, {"name": "right"}]
                },
                "network": {
                    "regions": [
                        {
                            "name": "vision",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "tensor", "values": [[0.1, 0.2], [0.3, 0.4]]},
            }
        ),
        encoding="utf-8",
    )

    with pytest.raises(ValueError, match="exactly one configured port"):
        load_experiment_config_file(str(config_path))


def test_cli_override_applies_external_input_path_and_runtime_defaults(tmp_path):
    override_input_path = tmp_path / "override.json"
    override_input_path.write_text(json.dumps([0.3, 0.6]), encoding="utf-8")

    config_path = tmp_path / "override_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "override_probe",
                "runtime": {"ticks": 1, "trace": True},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.1]},
            }
        ),
        encoding="utf-8",
    )

    experiment_config = load_experiment_config_file(str(config_path))
    experiment_config = apply_cli_overrides(experiment_config, input_path_override=str(override_input_path))
    trace_config = TraceConfig.parse("python", ["--lang", "python", "--config", str(config_path), "--ticks", "2"])
    trace_config = apply_trace_config_defaults(trace_config, experiment_config)

    loaded_schedule = load_input_schedule(experiment_config)

    assert trace_config.ticks == 2
    assert trace_config.trace_enabled is True
    assert loaded_schedule.tick_payloads == [{"a": 0.3}, {"a": 0.6}]


def test_batch_run_from_config_executes_and_writes_metrics_file(tmp_path):
    config_path = tmp_path / "run_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "run_probe",
                "runtime": {"ticks": 2},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.2, 0.4]},
                "output": {"metrics_file": "metrics.json"},
            }
        ),
        encoding="utf-8",
    )

    output = run_for_test(["--lang", "python", "--config", str(config_path)])
    metrics_path = tmp_path / "metrics.json"

    assert "experiment=run_probe tick=1" in output
    assert "experiment=run_probe tick=2" in output
    assert metrics_path.exists()
    metrics_payload = json.loads(metrics_path.read_text(encoding="utf-8"))
    assert metrics_payload["experiment"] == "run_probe"


def test_output_directory_override_writes_default_metrics_file(tmp_path):
    config_path = tmp_path / "directory_output_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "directory_output_probe",
                "runtime": {"ticks": 1},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.4]},
            }
        ),
        encoding="utf-8",
    )

    output_directory = tmp_path / "artifacts"
    run_for_test(["--lang", "python", "--config", str(config_path), "--output", str(output_directory)])

    metrics_path = output_directory / "metrics.json"
    assert metrics_path.exists()


def test_trace_with_config_emits_tick_events_as_jsonl(tmp_path):
    config_path = tmp_path / "trace_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "trace_probe",
                "runtime": {"ticks": 1},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.4]},
            }
        ),
        encoding="utf-8",
    )

    output = run_for_test(
        [
            "--lang",
            "python",
            "--trace",
            "--trace-format",
            "jsonl",
            "--config",
            str(config_path),
        ]
    )
    payloads = [json.loads(line_value) for line_value in output.splitlines() if line_value.strip()]

    assert payloads[0]["event"] == "tick.start"
    assert payloads[-1]["event"] == "tick.end"
    assert payloads[0]["scenario"] == "trace_probe"


def test_benchmark_with_config_emits_summary_lines(tmp_path):
    config_path = tmp_path / "benchmark_config.json"
    config_path.write_text(
        json.dumps(
            {
                "name": "benchmark_probe",
                "runtime": {"ticks": 1},
                "input_space": {"rank": 1, "shape": [1], "ports": [{"name": "a"}]},
                "network": {
                    "regions": [
                        {
                            "name": "main",
                            "layers": [
                                {"excitatory_count": 1, "inhibitory_count": 0, "modulatory_count": 0}
                            ]
                        }
                    ]
                },
                "input": {"type": "sequence", "values": [0.4]},
            }
        ),
        encoding="utf-8",
    )

    output = run_for_test(
        [
            "--lang",
            "python",
            "--benchmark",
            "--config",
            str(config_path),
            "--benchmark-warmup",
            "1",
            "--benchmark-iterations",
            "2",
            "--benchmark-repeat",
            "1",
        ]
    )

    assert "[benchmark.start]" in output
    assert "scenario=benchmark_probe" in output
    assert "[benchmark.summary]" in output
