from __future__ import annotations

import copy
from dataclasses import dataclass, field

from debug import trace as grownet_trace
from debug.trace_config import TraceConfig
from metrics import RegionMetrics
from region import Region
from runner.experiment_config import InputSpaceConfig
from runner.scaffold_builder import build_region_from_scaffold

from .interactive_keyboard import InteractiveCommand, interpret_key, map_digit_to_value
from .interactive_profiles import InteractiveProfile, RegionScaffold


@dataclass(frozen=True)
class InteractiveSnapshot:
    event: str
    tick: int
    profile: str
    input_values: dict[str, float]
    delivered_events: int
    total_slots: int
    total_synapses: int
    neuron_count: int
    layer_count: int
    region_count: int
    fired: bool
    trigger_key: str | None = None
    updated_port: str | None = None
    updated_value: float | None = None
    slot_delta: int = 0
    neuron_delta: int = 0
    layer_delta: int = 0
    trace_lines: tuple[str, ...] = ()
    focus_policy: str | None = None
    focus_point: tuple[int, int] | None = None
    focus_value: float | None = None
    anchor_point: tuple[int, int] | None = None
    anchor_value: float | None = None
    anchor_count: int = 0
    focus_candidate_count: int = 0


@dataclass(frozen=True)
class InteractiveSessionEvent:
    kind: str
    snapshot: InteractiveSnapshot | None = None
    message: str | None = None
    should_quit: bool = False
    should_reprint_panel: bool = False


@dataclass
class NetworkSnapshot:
    total_slots: int = 0
    total_synapses: int = 0
    neuron_count: int = 0
    layer_count: int = 0
    region_count: int = 0
    fired: bool = False


class InteractiveSession:
    def __init__(
        self,
        profile: InteractiveProfile,
        normalize_inputs: str | None = None,
        output_mode: str | None = None,
        trace_enabled: bool | None = None,
    ) -> None:
        self.original_profile = copy.deepcopy(profile)
        self.profile = copy.deepcopy(profile)
        self.normalize_inputs = normalize_inputs or self.profile.default_normalize_inputs
        self.output_mode = output_mode or self.profile.default_output
        self.trace_enabled = self.profile.default_trace if trace_enabled is None else bool(trace_enabled)
        self.tick_number = 0
        self.input_values: dict[str, float] = {}
        self.regions: list[Region] = []
        self.last_metrics = RegionMetrics()
        self.last_snapshot = NetworkSnapshot()
        self.trace_capture_writer = TraceCaptureWriter()
        self.reset()

    def reset(self) -> None:
        self.profile = copy.deepcopy(self.original_profile)
        self.tick_number = 0
        self.input_values = {port_name: 0.0 for port_name in self.profile.input_ports}
        self.regions = [self._build_region(region_scaffold) for region_scaffold in self.profile.regions]
        self.last_metrics = RegionMetrics()
        self.last_snapshot = self._snapshot_network()
        grownet_trace.reset()
        if self.trace_enabled:
            self._configure_trace_capture()

    def handle_key(self, raw_key: str) -> InteractiveSessionEvent:
        command = interpret_key(raw_key, self.profile.input_ports)
        if command.action == "input_tick":
            return self._handle_input_tick(command)
        if command.action == "repeat_tick":
            return self._run_tick(trigger_key=raw_key, updated_port=None, updated_value=None)
        if command.action == "summary":
            return InteractiveSessionEvent(kind="snapshot", snapshot=self.build_summary_snapshot())
        if command.action == "reset":
            self.reset()
            return InteractiveSessionEvent(
                kind="message",
                message="Session reset to the original scaffold.",
                should_reprint_panel=True,
            )
        if command.action == "toggle_output":
            return self._toggle_output_mode()
        if command.action == "clear":
            return InteractiveSessionEvent(kind="panel", should_reprint_panel=True)
        if command.action == "help":
            return InteractiveSessionEvent(kind="help")
        if command.action == "quit":
            return InteractiveSessionEvent(kind="quit", should_quit=True)
        return InteractiveSessionEvent(kind="message", message=f"Ignored key: {raw_key!r}")

    def build_summary_snapshot(self) -> InteractiveSnapshot:
        current_snapshot = self._snapshot_network()
        return InteractiveSnapshot(
            event="interactive.summary",
            tick=self.tick_number,
            profile=self.profile.name,
            input_values=dict(self.input_values),
            delivered_events=int(getattr(self.last_metrics, "delivered_events", 0)),
            total_slots=current_snapshot.total_slots,
            total_synapses=current_snapshot.total_synapses,
            neuron_count=current_snapshot.neuron_count,
            layer_count=current_snapshot.layer_count,
            region_count=current_snapshot.region_count,
            fired=current_snapshot.fired,
            **self._focus_snapshot_fields(),
        )

    def _handle_input_tick(self, command: InteractiveCommand) -> InteractiveSessionEvent:
        if command.updated_port is None or command.digit_value is None:
            return InteractiveSessionEvent(kind="message", message=f"Ignored key: {command.raw_key!r}")
        mapped_value = map_digit_to_value(command.digit_value, self.normalize_inputs)
        self.input_values[command.updated_port] = mapped_value
        return self._run_tick(
            trigger_key=command.raw_key,
            updated_port=command.updated_port,
            updated_value=mapped_value,
        )

    def _run_tick(
        self,
        trigger_key: str | None,
        updated_port: str | None,
        updated_value: float | None,
    ) -> InteractiveSessionEvent:
        before_snapshot = self._snapshot_network()
        if self.trace_enabled:
            self.trace_capture_writer.reset()
            grownet_trace.begin_tick(
                self.profile.name,
                updated_port or "interactive",
                updated_value,
                use_default_signal_id=True,
            )

        aggregate_metrics = RegionMetrics()
        for region_obj in self.regions:
            ordered_port_values: dict[str, float] = {}
            for port_name in self.profile.input_ports:
                if port_name in region_obj.input_edges:
                    ordered_port_values[port_name] = float(self.input_values[port_name])
            region_metrics = region_obj.tick_inputs(ordered_port_values)
            aggregate_metrics.inc_delivered_events(region_metrics.delivered_events)
            aggregate_metrics.add_slots(region_metrics.total_slots)
            aggregate_metrics.add_synapses(region_metrics.total_synapses)

        if self.trace_enabled:
            grownet_trace.end_tick(aggregate_metrics)

        self.tick_number += 1
        self.last_metrics = aggregate_metrics
        after_snapshot = self._snapshot_network()
        self.last_snapshot = after_snapshot
        snapshot = InteractiveSnapshot(
            event="interactive.tick",
            tick=self.tick_number,
            profile=self.profile.name,
            trigger_key=trigger_key,
            updated_port=updated_port,
            updated_value=updated_value,
            input_values=dict(self.input_values),
            delivered_events=int(aggregate_metrics.delivered_events),
            total_slots=after_snapshot.total_slots,
            total_synapses=after_snapshot.total_synapses,
            neuron_count=after_snapshot.neuron_count,
            layer_count=after_snapshot.layer_count,
            region_count=after_snapshot.region_count,
            fired=after_snapshot.fired,
            slot_delta=after_snapshot.total_slots - before_snapshot.total_slots,
            neuron_delta=after_snapshot.neuron_count - before_snapshot.neuron_count,
            layer_delta=after_snapshot.layer_count - before_snapshot.layer_count,
            trace_lines=tuple(self.trace_capture_writer.lines if self.trace_enabled else []),
            **self._focus_snapshot_fields(),
        )
        return InteractiveSessionEvent(kind="snapshot", snapshot=snapshot)

    def _toggle_output_mode(self) -> InteractiveSessionEvent:
        if self.output_mode == "jsonl":
            return InteractiveSessionEvent(
                kind="message",
                message="Output mode remains jsonl; use restart flags to switch from jsonl.",
            )
        self.output_mode = "verbose" if self.output_mode == "compact" else "compact"
        return InteractiveSessionEvent(kind="message", message=f"Output mode set to {self.output_mode}.")

    def _build_region(self, region_scaffold: RegionScaffold) -> Region:
        input_space = InputSpaceConfig(
            rank=int(self.profile.rank),
            shape=(len(self.profile.input_ports),) if int(self.profile.rank) == 1 else (1,) * int(self.profile.rank),
            ports=tuple(self.profile.input_ports),
        )
        return build_region_from_scaffold(region_scaffold, input_space, focus_config=self.profile.focus_config)

    def _configure_trace_capture(self) -> None:
        trace_format = "jsonl" if self.output_mode == "jsonl" else "text"
        trace_config = TraceConfig(
            trace_enabled=True,
            lang="python",
            trace_format=trace_format,
            scenario=self.profile.name,
            value=0.0,
            ticks=1,
        )
        self.trace_capture_writer.reset()
        grownet_trace.configure(trace_config, self.trace_capture_writer)

    def _snapshot_network(self) -> NetworkSnapshot:
        snapshot = NetworkSnapshot(region_count=len(self.regions))
        for region_obj in self.regions:
            region_layers = region_obj.get_layers()
            snapshot.layer_count += len(region_layers)
            for layer_obj in region_layers:
                neurons = layer_obj.get_neurons()
                snapshot.neuron_count += len(neurons)
                for neuron_obj in neurons:
                    slots_map = getattr(neuron_obj, "slots", {})
                    outgoing_list = neuron_obj.get_outgoing() if hasattr(neuron_obj, "get_outgoing") else []
                    snapshot.total_slots += len(slots_map) if isinstance(slots_map, dict) else 0
                    snapshot.total_synapses += len(outgoing_list)
                    snapshot.fired = snapshot.fired or bool(getattr(neuron_obj, "fired_last", False))
        return snapshot

    def _focus_snapshot_fields(self) -> dict[str, object]:
        for region_obj in self.regions:
            focus_summary = region_obj.get_focus_summary()
            if focus_summary is None:
                continue
            return {
                "focus_policy": getattr(focus_summary, "policy_name", None),
                "focus_point": getattr(focus_summary, "active_focus_point", None),
                "focus_value": getattr(focus_summary, "active_focus_value", None),
                "anchor_point": getattr(focus_summary, "active_anchor_point", None),
                "anchor_value": getattr(focus_summary, "active_anchor_value", None),
                "anchor_count": int(getattr(focus_summary, "anchor_count", 0)),
                "focus_candidate_count": int(getattr(focus_summary, "candidate_count", 0)),
            }
        return {
            "focus_policy": None,
            "focus_point": None,
            "focus_value": None,
            "anchor_point": None,
            "anchor_value": None,
            "anchor_count": 0,
            "focus_candidate_count": 0,
        }


class TraceCaptureWriter:
    def __init__(self) -> None:
        self.lines: list[str] = []
        self._buffer = ""

    def reset(self) -> None:
        self.lines = []
        self._buffer = ""

    def write(self, payload: str) -> int:
        self._buffer += str(payload)
        while "\n" in self._buffer:
            line_value, remainder = self._buffer.split("\n", 1)
            if line_value:
                self.lines.append(line_value)
            self._buffer = remainder
        return len(payload)

    def flush(self) -> None:
        if self._buffer:
            self.lines.append(self._buffer)
            self._buffer = ""

    def close(self) -> None:
        self.flush()
