from __future__ import annotations

import sys
from dataclasses import replace
from dataclasses import dataclass
from pathlib import Path

THIS_FILE = Path(__file__).resolve()
PYTHON_ROOT = THIS_FILE.parent.parent
if str(PYTHON_ROOT) not in sys.path:
    sys.path.insert(0, str(PYTHON_ROOT))

from debug import trace as grownet_trace
from debug.benchmark_runner import run as run_benchmark
from debug.trace_config import TraceConfig
from growth import GrowthPolicy, maybe_grow
from metrics import RegionMetrics
from region import Region
from runner.experiment_loader import (
    apply_cli_overrides,
    apply_trace_config_defaults,
    load_experiment_config_file,
)
from runner.run_mode_dispatch import execute_experiment


@dataclass(frozen=True)
class TraceRunSummary:
    delivered_events: int
    total_slots: int
    total_synapses: int
    layer_count: int
    neuron_count: int


def main(argv: list[str] | None = None) -> int:
    args = list(sys.argv[1:] if argv is None else argv)
    trace_config = TraceConfig.parse("python", args)
    if trace_config.help_requested:
        print(TraceConfig.usage("python"))
        return 0
    trace_config, experiment_config = _prepare_execution(trace_config)
    if trace_config.mode == "interactive":
        print("GrowNet experiment runner: use --interactive for Python interactive mode.", file=sys.stderr)
        return 2

    if trace_config.benchmark_enabled:
        benchmark_writer = _open_benchmark_writer(trace_config)
        trace_writer = _open_benchmark_trace_writer(trace_config)
        try:
            run_benchmark(
                trace_config,
                benchmark_writer,
                trace_writer,
                lambda benchmark_trace_config, benchmark_trace_writer: execute_once(
                    benchmark_trace_config,
                    benchmark_trace_writer,
                    experiment_config=experiment_config,
                ),
            )
        finally:
            _close_writer(trace_config.writes_benchmark_to_stdout, benchmark_writer)
            _close_writer(trace_config.writes_to_stdout, trace_writer)
        return 0

    trace_writer = _open_writer(trace_config)
    try:
        run(trace_config, trace_writer, experiment_config=experiment_config)
    finally:
        _close_writer(trace_config.writes_to_stdout, trace_writer)
    return 0


def run(trace_config: TraceConfig, writer, experiment_config=None) -> None:
    _execute(trace_config, writer, emit_plain_summary=True, experiment_config=experiment_config)


def execute_once(trace_config: TraceConfig, writer, experiment_config=None) -> TraceRunSummary:
    return _execute(trace_config, writer, emit_plain_summary=False, experiment_config=experiment_config)


def run_for_test(argv: list[str]) -> str:
    from io import StringIO

    trace_config = TraceConfig.parse("python", argv)
    trace_config, experiment_config = _prepare_execution(trace_config)
    output_buffer = StringIO()
    if trace_config.benchmark_enabled:
        run_benchmark(
            trace_config,
            output_buffer,
            _DiscardingWriter(),
            lambda benchmark_trace_config, benchmark_trace_writer: execute_once(
                benchmark_trace_config,
                benchmark_trace_writer,
                experiment_config=experiment_config,
            ),
        )
    else:
        _execute(trace_config, output_buffer, emit_plain_summary=True, experiment_config=experiment_config)
    return output_buffer.getvalue()


def run_summary_for_test(argv: list[str], writer=None) -> TraceRunSummary:
    from io import StringIO

    trace_config = TraceConfig.parse("python", argv)
    trace_config, experiment_config = _prepare_execution(trace_config)
    output_writer = writer if writer is not None else StringIO()
    return _execute(
        trace_config,
        output_writer,
        emit_plain_summary=False,
        experiment_config=experiment_config,
    )


def _open_writer(trace_config: TraceConfig):
    if trace_config.writes_to_stdout:
        return sys.stdout
    trace_config.out_path.parent.mkdir(parents=True, exist_ok=True)
    return trace_config.out_path.open("w", encoding="utf-8")


def _open_benchmark_writer(trace_config: TraceConfig):
    if trace_config.writes_benchmark_to_stdout:
        return sys.stdout
    trace_config.benchmark_out_path.parent.mkdir(parents=True, exist_ok=True)
    return trace_config.benchmark_out_path.open("w", encoding="utf-8")


def _open_benchmark_trace_writer(trace_config: TraceConfig):
    if not trace_config.trace_enabled or trace_config.writes_to_stdout:
        return _DiscardingWriter()
    return _open_writer(trace_config)


def _close_writer(writes_to_stdout: bool, writer) -> None:
    if writes_to_stdout:
        writer.flush()
        return
    writer.close()


def _build_single_neuron_region() -> Region:
    region = Region("dbg")
    region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    return region


def _build_two_layer_region(value: float) -> Region:
    region = Region("dbg")
    source_layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    dest_layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    region.connect_layers(source_layer_index, dest_layer_index, probability=1.0, feedback=False)
    _prime_two_layer_source(region.get_layers()[source_layer_index].get_neurons()[0], value)
    return region


def _build_image_region() -> Region:
    region = Region("dbg")
    layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    neuron = region.get_layers()[layer_index].get_neurons()[0]
    neuron.slot_cfg.spatial_enabled = True
    neuron.slot_cfg.row_bin_width_pct = 25.0
    neuron.slot_cfg.col_bin_width_pct = 25.0
    neuron.slot_limit = 16
    return region


def _build_growth_region(value: float) -> Region:
    region = Region("dbg")
    layer_index = region.add_layer(excitatory_count=1, inhibitory_count=0, modulatory_count=0)
    neuron = region.get_layers()[layer_index].get_neurons()[0]
    neuron.slot_cfg.slot_limit = 1
    neuron.slot_limit = 1
    neuron.slot_cfg.fallback_growth_threshold = 1
    neuron.slot_cfg.neuron_growth_cooldown_ticks = 0
    neuron.slot_cfg.growth_enabled = True
    neuron.slot_cfg.neuron_growth_enabled = True
    neuron.slot_cfg.layer_growth_enabled = False
    region.set_growth_policy(
        GrowthPolicy(
            enable_layer_growth=True,
            max_total_layers=16,
            avg_slots_threshold=0.0,
            percent_neurons_at_cap_threshold=0.0,
            layer_cooldown_ticks=0,
            new_layer_excitatory_count=2,
            wire_probability=1.0,
        )
    )
    baseline_value = 0.1 if abs(value) < 1e-6 else value / 2.0
    neuron.on_input(baseline_value)
    _finish_tick(region, allow_region_growth=False)
    return region


def _run_single_neuron_tick(region: Region, value: float) -> RegionMetrics:
    layer_obj = region.get_layers()[0]
    grownet_trace.begin_tick(region.get_name(), "x", value, use_default_signal_id=True)
    grownet_trace.emit_layer_forward(layer_obj, value)
    layer_obj.forward(value)
    return _finish_tick(region, allow_region_growth=False)


def _run_two_layer_tick(region: Region, value: float) -> RegionMetrics:
    layer_obj = region.get_layers()[0]
    grownet_trace.begin_tick(region.get_name(), "x", value, use_default_signal_id=True)
    grownet_trace.emit_layer_forward(layer_obj, value)
    layer_obj.forward(value)
    return _finish_tick(region, allow_region_growth=False)


def _run_image_tick(region: Region, value: float) -> RegionMetrics:
    layer_obj = region.get_layers()[0]
    tick_number = grownet_trace.begin_tick(region.get_name(), "img", None, use_default_signal_id=False)
    frame = [
        [value, 0.0],
        [0.0, value / 2.0],
    ]
    for row_index, row_values in enumerate(frame):
        for col_index, pixel_value in enumerate(row_values):
            if pixel_value == 0.0:
                continue
            signal_id = grownet_trace.format_pixel_signal_id(tick_number, row_index, col_index)
            with grownet_trace.run_with_signal(signal_id):
                event_fields: dict[str, object] = {}
                grownet_trace.add_layer_fields(event_fields, layer_obj)
                event_fields["neuron_count"] = len(layer_obj.get_neurons())
                event_fields["input_value"] = float(pixel_value)
                event_fields["reason"] = "pixel"
                grownet_trace.emit("layer.forward", **event_fields)
                neuron = layer_obj.get_neurons()[0]
                fired = neuron.on_input_2d(pixel_value, row_index, col_index)
                if fired:
                    neuron.on_output(pixel_value)
    return _finish_tick(region, allow_region_growth=False)


def _run_growth_tick(region: Region, value: float) -> RegionMetrics:
    layer_obj = region.get_layers()[0]
    grownet_trace.begin_tick(region.get_name(), "x", value, use_default_signal_id=True)
    grownet_trace.emit_layer_forward(layer_obj, value)
    layer_obj.forward(value)
    return _finish_tick(region, allow_region_growth=True)


def _finish_tick(region: Region, allow_region_growth: bool) -> RegionMetrics:
    for layer_obj in region.get_layers():
        layer_obj.end_tick()
        grownet_trace.emit_bus_decay(layer_obj)
    if allow_region_growth and region.get_growth_policy() is not None:
        maybe_grow(region, region.get_growth_policy())
    metrics = _collect_metrics(region)
    grownet_trace.end_tick(metrics)
    return metrics


def _collect_metrics(region: Region) -> RegionMetrics:
    metrics = RegionMetrics()
    metrics.inc_delivered_events(1)
    for layer_obj in region.get_layers():
        for neuron in layer_obj.get_neurons():
            metrics.add_slots(len(getattr(neuron, "slots", {})))
            metrics.add_synapses(len(getattr(neuron, "outgoing", [])))
    return metrics


def _execute(trace_config: TraceConfig, writer, emit_plain_summary: bool, experiment_config=None) -> TraceRunSummary:
    if experiment_config is not None:
        return _execute_config_experiment(
            trace_config,
            writer,
            emit_plain_summary=emit_plain_summary,
            experiment_config=experiment_config,
        )

    grownet_trace.configure(trace_config, writer)
    try:
        region = {
            "single_neuron": _build_single_neuron_region,
            "two_layer": lambda: _build_two_layer_region(trace_config.value),
            "image_2d": _build_image_region,
            "growth": lambda: _build_growth_region(trace_config.value),
        }[trace_config.scenario]()

        metrics = RegionMetrics()
        for tick_ordinal in range(trace_config.ticks):
            if trace_config.scenario == "single_neuron":
                metrics = _run_single_neuron_tick(region, trace_config.value)
            elif trace_config.scenario == "two_layer":
                metrics = _run_two_layer_tick(region, trace_config.value)
            elif trace_config.scenario == "image_2d":
                metrics = _run_image_tick(region, trace_config.value)
            elif trace_config.scenario == "growth":
                metrics = _run_growth_tick(region, trace_config.value)
            else:
                raise ValueError(f"Unsupported scenario: {trace_config.scenario}")

            if emit_plain_summary and not trace_config.trace_enabled:
                writer.write(
                    "scenario={scenario} tick={tick} delivered_events={delivered_events} "
                    "total_slots={total_slots} total_synapses={total_synapses}\n".format(
                        scenario=trace_config.scenario,
                        tick=tick_ordinal + 1,
                        delivered_events=metrics.delivered_events,
                        total_slots=metrics.total_slots,
                        total_synapses=metrics.total_synapses,
                    )
                )
                writer.flush()
        return TraceRunSummary(
            delivered_events=int(metrics.delivered_events),
            total_slots=int(metrics.total_slots),
            total_synapses=int(metrics.total_synapses),
            layer_count=len(region.get_layers()),
            neuron_count=sum(len(layer_obj.get_neurons()) for layer_obj in region.get_layers()),
        )
    finally:
        grownet_trace.reset()


def _execute_config_experiment(trace_config: TraceConfig, writer, emit_plain_summary: bool, experiment_config) -> TraceRunSummary:
    config_trace_config = replace(trace_config, scenario=experiment_config.name)
    grownet_trace.configure(config_trace_config, writer)
    try:
        experiment_summary = execute_experiment(
            experiment_config,
            config_trace_config,
            writer,
            emit_plain_summary=emit_plain_summary,
        )
        return TraceRunSummary(
            delivered_events=int(experiment_summary.delivered_events),
            total_slots=int(experiment_summary.total_slots),
            total_synapses=int(experiment_summary.total_synapses),
            layer_count=int(experiment_summary.layer_count),
            neuron_count=int(experiment_summary.neuron_count),
        )
    finally:
        grownet_trace.reset()


def _prepare_execution(trace_config: TraceConfig):
    if trace_config.config_path is None:
        return trace_config, None
    experiment_config = load_experiment_config_file(trace_config.config_path)
    experiment_config = apply_cli_overrides(
        experiment_config,
        input_path_override=trace_config.input_path_override,
        input_type_override=trace_config.input_type_override,
        output_override=None,
    )
    effective_trace_config = apply_trace_config_defaults(trace_config, experiment_config)
    if (
        effective_trace_config.generic_output_target is not None
        and not effective_trace_config.trace_enabled
        and not effective_trace_config.benchmark_enabled
    ):
        experiment_config = apply_cli_overrides(
            experiment_config,
            output_override=effective_trace_config.generic_output_target,
        )
    return effective_trace_config, experiment_config


def _prime_two_layer_source(neuron, value: float) -> None:
    neuron.on_input(value)
    slot_obj = neuron.slots.get(0)
    if slot_obj is None:
        return
    slot_obj.set_strength_value(max(0.9, abs(value) + 0.2))
    slot_obj.set_threshold_value(max(0.0, abs(value) * 0.25))
    slot_obj.set_hit_count(100)


class _DiscardingWriter:
    def write(self, payload: str) -> int:
        return len(payload)

    def flush(self) -> None:
        return None

    def close(self) -> None:
        return None


if __name__ == "__main__":
    raise SystemExit(main())
