from __future__ import annotations

from dataclasses import dataclass, field
from math import prod
from typing import Any, Callable

from focus import FocusConfig, normalize_focus_policy_name

VALID_INPUT_TYPES = {"sequence", "tensor", "image", "audio"}
VALID_RUNTIME_MODES = {"run", "trace", "benchmark", "interactive"}


@dataclass(frozen=True)
class LayerScaffold:
    excitatory_count: int
    inhibitory_count: int
    modulatory_count: int
    neuron_limit: int | None = None


@dataclass(frozen=True)
class ConnectionScaffold:
    source_layer_index: int
    dest_layer_index: int
    probability: float = 1.0
    feedback: bool = False


@dataclass(frozen=True)
class RegionScaffold:
    name: str
    layers: list[LayerScaffold]
    input_bindings: dict[str, list[int]] = field(default_factory=dict)
    connections: list[ConnectionScaffold] = field(default_factory=list)
    growth_policy: dict[str, Any] | None = None
    post_build_hook: Callable[[object], None] | None = None


@dataclass(frozen=True)
class RuntimeDefaultsConfig:
    ticks: int | None = None
    trace: bool = False
    trace_format: str = "text"
    benchmark: bool = False
    interactive: bool = False


@dataclass(frozen=True)
class InteractiveDefaultsConfig:
    output: str = "compact"
    trace: bool = False
    normalize_inputs: str = "normalized"


@dataclass(frozen=True)
class InputSpaceConfig:
    rank: int
    shape: tuple[int, ...]
    ports: tuple[str, ...]
    gain: float = 1.0
    epsilon_fire: float = 0.01

    def total_size(self) -> int:
        return prod(self.shape) if self.shape else 0

    def uses_scalar_port_inputs(self) -> bool:
        return self.rank == 1 and self.total_size() == len(self.ports) and len(self.ports) > 0

    def primary_port(self) -> str:
        if not self.ports:
            raise ValueError("Experiment input space must define at least one port.")
        return self.ports[0]


@dataclass(frozen=True)
class InputSourceConfig:
    input_type: str
    path: Any = None
    values: Any = None
    mode: str = "grayscale"
    resize: tuple[int, ...] | None = None
    sample_rate: int | None = None
    frame_size: int | None = None
    hop_size: int | None = None


@dataclass(frozen=True)
class OutputConfig:
    trace_file: Any = None
    metrics_file: Any = None
    benchmark_file: Any = None


@dataclass(frozen=True)
class ExperimentConfig:
    name: str
    runtime: RuntimeDefaultsConfig
    interactive: InteractiveDefaultsConfig
    focus: FocusConfig
    input_space: InputSpaceConfig
    regions: list[RegionScaffold]
    input_source: InputSourceConfig
    output: OutputConfig
    source_path: Any = None

    def default_mode(self) -> str:
        if self.runtime.interactive:
            return "interactive"
        if self.runtime.benchmark:
            return "benchmark"
        if self.runtime.trace:
            return "trace"
        return "run"


def validate_runtime_mode(runtime_mode: str | None) -> str:
    normalized_mode = str(runtime_mode or "run").strip().lower()
    if normalized_mode not in VALID_RUNTIME_MODES:
        raise ValueError(f"Unsupported runtime mode: {runtime_mode}")
    return normalized_mode


def parse_port_names(ports_payload: list[dict[str, Any]] | list[str] | None) -> list[str]:
    if not ports_payload:
        return ["a"]

    port_names: list[str] = []
    for port_payload in ports_payload:
        if isinstance(port_payload, dict):
            port_name = str(port_payload.get("name", "")).strip()
        else:
            port_name = str(port_payload).strip()
        if not port_name:
            raise ValueError("Input ports must have non-empty names.")
        port_names.append(port_name)

    if not port_names:
        raise ValueError("At least one input port is required.")
    return port_names


def parse_layer_scaffold(layer_payload: dict[str, Any]) -> LayerScaffold:
    return LayerScaffold(
        excitatory_count=int(layer_payload.get("excitatory_count", 0)),
        inhibitory_count=int(layer_payload.get("inhibitory_count", 0)),
        modulatory_count=int(layer_payload.get("modulatory_count", 0)),
        neuron_limit=(
            int(layer_payload["neuron_limit"])
            if layer_payload.get("neuron_limit") is not None
            else None
        ),
    )


def parse_input_bindings(
    bindings_payload: dict[str, Any] | None,
    input_ports: list[str],
    layer_count: int,
) -> dict[str, list[int]]:
    if not bindings_payload:
        return {port_name: [0] for port_name in input_ports}

    parsed_bindings: dict[str, list[int]] = {}
    for port_name in input_ports:
        raw_indices = bindings_payload.get(port_name, [0])
        if not isinstance(raw_indices, list):
            raise ValueError(f"Input binding for port '{port_name}' must be a list of layer indices.")
        parsed_indices = [int(layer_index) for layer_index in raw_indices]
        for layer_index in parsed_indices:
            if layer_index < 0 or layer_index >= layer_count:
                raise ValueError(f"Input binding out of range: {port_name} -> {layer_index}")
        parsed_bindings[port_name] = parsed_indices
    return parsed_bindings


def parse_connections(
    connections_payload: list[dict[str, Any]] | None,
    layer_count: int,
) -> list[ConnectionScaffold]:
    if not connections_payload:
        return []

    parsed_connections: list[ConnectionScaffold] = []
    for connection_payload in connections_payload:
        source_layer_index = int(connection_payload.get("source_layer_index", -1))
        dest_layer_index = int(connection_payload.get("dest_layer_index", -1))
        if source_layer_index < 0 or source_layer_index >= layer_count:
            raise ValueError(f"Connection source layer out of range: {source_layer_index}")
        if dest_layer_index < 0 or dest_layer_index >= layer_count:
            raise ValueError(f"Connection destination layer out of range: {dest_layer_index}")
        parsed_connections.append(
            ConnectionScaffold(
                source_layer_index=source_layer_index,
                dest_layer_index=dest_layer_index,
                probability=float(connection_payload.get("probability", 1.0)),
                feedback=bool(connection_payload.get("feedback", False)),
            )
        )
    return parsed_connections


def parse_region_scaffold(region_payload: dict[str, Any], input_ports: list[str]) -> RegionScaffold:
    layer_payloads = region_payload.get("layers", [])
    if not layer_payloads:
        raise ValueError("Each configured region must define at least one layer.")

    layers = [parse_layer_scaffold(layer_payload) for layer_payload in layer_payloads]
    input_bindings = parse_input_bindings(region_payload.get("input_bindings"), input_ports, len(layers))
    connections = parse_connections(region_payload.get("connections"), len(layers))

    if not connections and len(layers) > 1:
        connections = [
            ConnectionScaffold(
                source_layer_index=layer_index,
                dest_layer_index=layer_index + 1,
                probability=1.0,
                feedback=False,
            )
            for layer_index in range(len(layers) - 1)
        ]

    growth_policy_payload = region_payload.get("growth_policy")
    growth_policy = dict(growth_policy_payload) if isinstance(growth_policy_payload, dict) else None

    return RegionScaffold(
        name=str(region_payload.get("name", "main")),
        layers=layers,
        input_bindings=input_bindings,
        connections=connections,
        growth_policy=growth_policy,
    )


def parse_focus_config(focus_payload: dict[str, Any] | None) -> FocusConfig:
    payload = dict(focus_payload or {})
    return FocusConfig(
        enabled=bool(payload.get("enabled", False)),
        policy=normalize_focus_policy_name(payload.get("policy", "energy_first")),
        serial=bool(payload.get("serial", True)),
        candidate_limit=max(1, int(payload.get("candidate_limit", 8))),
        allow_anchor_map=bool(payload.get("allow_anchor_map", True)),
        inhibition_of_return_ticks=max(0, int(payload.get("inhibition_of_return_ticks", 3))),
        anchor_map_capacity=max(1, int(payload.get("anchor_map_capacity", 8))),
        anchor_decay_ticks=max(1, int(payload.get("anchor_decay_ticks", 64))),
    )
