from __future__ import annotations

import functools
import json
import shutil
import subprocess
from pathlib import Path

import debug.trace as grownet_trace
from demos.grownet_trace_main import run_for_test, run_summary_for_test
from growth import GrowthPolicy, maybe_grow
from neuron import Neuron
from region import Region
from slot_config import fixed
from slot_engine import pack_slot_key_2d, slot_key_to_trace_id
from weight import Weight


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


def test_single_neuron_summary_matches_with_trace_on_and_off():
    assert_trace_on_off_equivalent("single_neuron", 1)


def test_two_layer_summary_matches_with_trace_on_and_off():
    assert_trace_on_off_equivalent("two_layer", 1)


def test_growth_summary_matches_with_trace_on_and_off():
    assert_trace_on_off_equivalent("growth", 1)


def test_failing_trace_writer_does_not_abort_tick_paths():
    baseline_summary = run_summary_for_test(
        ["--lang", "python", "--trace", "--trace-format", "jsonl", "--scenario", "growth", "--value", "0.42"]
    )
    failing_writer_summary = run_summary_for_test(
        ["--lang", "python", "--trace", "--trace-format", "jsonl", "--scenario", "growth", "--value", "0.42"],
        writer=AlwaysFailingWriter(),
    )
    assert failing_writer_summary == baseline_summary


def test_formatter_failures_are_swallowed_locally(monkeypatch):
    baseline_summary = run_summary_for_test(
        ["--lang", "python", "--trace", "--trace-format", "jsonl", "--scenario", "single_neuron", "--value", "0.42"]
    )

    def raise_formatter_failure(event_fields):
        raise ValueError(f"synthetic formatter failure for {event_fields['event']}")

    monkeypatch.setattr(grownet_trace, "_format_json_line", raise_formatter_failure)
    failing_summary = run_summary_for_test(
        ["--lang", "python", "--trace", "--trace-format", "jsonl", "--scenario", "single_neuron", "--value", "0.42"]
    )
    assert failing_summary == baseline_summary


def test_request_layer_growth_defaults_preserve_previous_behavior():
    region = Region("layer_growth_defaults")
    saturated_layer_index = region.add_layer(excitatory_count=10, inhibitory_count=0, modulatory_count=0)
    saturated_layer = region.get_layers()[saturated_layer_index]

    new_layer_index = region.request_layer_growth(saturated_layer)

    new_layer = region.get_layers()[new_layer_index]
    assert len(new_layer.get_neurons()) == 5
    total_outgoing_edges = sum(len(neuron.get_outgoing()) for neuron in saturated_layer.get_neurons())
    assert total_outgoing_edges == 50


def test_policy_driven_growth_passes_connection_probability_and_excitatory_count():
    region = Region("policy_growth_forwarding")
    layer_index = region.add_layer(excitatory_count=2, inhibitory_count=0, modulatory_count=0)
    donor_layer = region.get_layers()[layer_index]
    for donor_neuron in donor_layer.get_neurons():
        donor_neuron.slot_limit = 1
        donor_neuron.slots[0] = Weight()
        donor_neuron.last_slot_used_fallback = True

    captured_growth_args: dict[str, float | int | None] = {}
    original_request_layer_growth = region.request_layer_growth

    def recording_request_layer_growth(saturated_layer, connection_probability=None, excitatory_count=None):
        captured_growth_args["connection_probability"] = connection_probability
        captured_growth_args["excitatory_count"] = excitatory_count
        return original_request_layer_growth(saturated_layer, connection_probability, excitatory_count)

    region.request_layer_growth = recording_request_layer_growth  # type: ignore[method-assign]
    growth_policy = GrowthPolicy(
        enable_layer_growth=True,
        max_total_layers=8,
        avg_slots_threshold=999.0,
        percent_neurons_at_cap_threshold=50.0,
        layer_cooldown_ticks=0,
        new_layer_excitatory_count=7,
        wire_probability=0.25,
    )

    assert maybe_grow(region, growth_policy) is True
    assert captured_growth_args == {"connection_probability": 0.25, "excitatory_count": 7}


def test_spatial_bookkeeping_uses_tuple_keys_while_trace_ids_stay_packed():
    spatial_config = fixed(10.0)
    spatial_config.spatial_enabled = True
    spatial_config.row_bin_width_pct = 25.0
    spatial_config.col_bin_width_pct = 25.0
    spatial_config.slot_limit = 2
    spatial_neuron = Neuron("spatial")
    spatial_neuron.slot_cfg = spatial_config
    spatial_neuron.slot_engine.cfg = spatial_config
    spatial_neuron.slot_limit = 2

    spatial_neuron.on_input_2d(1.0, 1, 1)
    spatial_neuron.on_input_2d(1.0, 2, 1)
    spatial_neuron.on_input_2d(1.0, 2, 2)

    assert isinstance(spatial_neuron.last_selected_slot_id, tuple)
    assert isinstance(spatial_neuron.last_missing_slot_id, tuple)
    assert spatial_neuron.last_selected_slot_id in spatial_neuron.slots
    assert slot_key_to_trace_id(spatial_neuron.last_selected_slot_id) == pack_slot_key_2d(spatial_neuron.last_selected_slot_id)
    assert slot_key_to_trace_id(spatial_neuron.last_missing_slot_id) == pack_slot_key_2d(spatial_neuron.last_missing_slot_id)


def test_java_python_single_neuron_event_schema_parity():
    python_payloads = parse_jsonl_output(
        run_for_test(
            ["--lang", "python", "--trace", "--trace-format", "jsonl", "--scenario", "single_neuron", "--value", "0.42"]
        )
    )
    java_payloads = parse_jsonl_output(run_java_trace_output("single_neuron"))

    assert [payload["event"] for payload in python_payloads] == [payload["event"] for payload in java_payloads]
    assert {payload["event"].split(".", 1)[0] for payload in python_payloads} == {
        payload["event"].split(".", 1)[0] for payload in java_payloads
    }

    for python_payload, java_payload in zip(python_payloads, java_payloads):
        assert python_payload["event"] == java_payload["event"]
        assert python_payload["tick"] == java_payload["tick"]
        assert python_payload["scenario"] == java_payload["scenario"]
        for field_name in ("signal_id", "layer_index", "neuron_id", "slot_id"):
            if field_name in python_payload or field_name in java_payload:
                assert python_payload.get(field_name) == java_payload.get(field_name)


def assert_trace_on_off_equivalent(scenario_name: str, tick_count: int):
    without_trace = run_summary_for_test(
        ["--lang", "python", "--scenario", scenario_name, "--ticks", str(tick_count), "--value", "0.42"]
    )
    with_trace = run_summary_for_test(
        [
            "--lang",
            "python",
            "--trace",
            "--trace-format",
            "jsonl",
            "--scenario",
            scenario_name,
            "--ticks",
            str(tick_count),
            "--value",
            "0.42",
        ]
    )
    assert with_trace == without_trace


def parse_jsonl_output(output_text: str) -> list[dict]:
    return [json.loads(line) for line in output_text.splitlines() if line.strip()]


@functools.lru_cache(maxsize=1)
def compile_java_trace_runner() -> None:
    maven_executable = shutil.which("mvn.cmd") or shutil.which("mvn")
    if maven_executable is None:
        raise FileNotFoundError("Could not find Maven on PATH.")
    subprocess.run(
        [maven_executable, "-q", "-DskipTests", "compile"],
        cwd=REPO_ROOT,
        check=True,
        capture_output=True,
        text=True,
    )


def run_java_trace_output(scenario_name: str) -> str:
    compile_java_trace_runner()
    completed_process = subprocess.run(
        [
            "java",
            "-cp",
            "target/classes",
            "ai.nektron.grownet.debug.GrowNetTraceMain",
            "--lang",
            "java",
            "--trace",
            "--trace-format",
            "jsonl",
            "--scenario",
            scenario_name,
            "--value",
            "0.42",
        ],
        cwd=REPO_ROOT,
        check=True,
        capture_output=True,
        text=True,
    )
    return completed_process.stdout


class AlwaysFailingWriter:
    def write(self, text: str) -> int:
        raise OSError(f"synthetic trace writer failure for {text!r}")

    def flush(self) -> None:
        raise OSError("synthetic trace writer failure")
