from __future__ import annotations

import json
from dataclasses import replace
from pathlib import Path
from typing import Any

from debug.trace_config import TraceConfig

from .experiment_config import (
    ExperimentConfig,
    InputSourceConfig,
    InputSpaceConfig,
    InteractiveDefaultsConfig,
    OutputConfig,
    RegionScaffold,
    RuntimeDefaultsConfig,
    parse_focus_config,
    parse_port_names,
    parse_region_scaffold,
)


def load_experiment_config_file(config_path: str) -> ExperimentConfig:
    config_file = Path(config_path).expanduser().resolve()
    if not config_file.exists():
        raise ValueError(f"Experiment config file not found: {config_path}")
    with config_file.open("r", encoding="utf-8") as config_handle:
        config_payload = json.load(config_handle)
    return parse_experiment_config(config_payload, config_file)


def parse_experiment_config(config_payload: dict[str, Any], source_path: Path | None = None) -> ExperimentConfig:
    base_directory = source_path.parent if source_path is not None else Path.cwd()
    input_space = _parse_input_space(config_payload.get("input_space"))
    runtime = _parse_runtime_defaults(config_payload.get("runtime", {}))
    interactive = _parse_interactive_defaults(config_payload.get("interactive", {}))
    focus = parse_focus_config(config_payload.get("focus", {}))
    regions = _parse_regions(config_payload.get("network", {}), list(input_space.ports))
    input_source = _parse_input_source(config_payload.get("input", {}), base_directory)
    output = _parse_output_config(config_payload.get("output", {}), base_directory)
    experiment_name = str(config_payload.get("name") or (source_path.stem if source_path is not None else "experiment"))
    return ExperimentConfig(
        name=experiment_name,
        runtime=runtime,
        interactive=interactive,
        focus=focus,
        input_space=input_space,
        regions=regions,
        input_source=input_source,
        output=output,
        source_path=source_path,
    )


def apply_cli_overrides(
    experiment_config: ExperimentConfig,
    input_path_override: str | None = None,
    input_type_override: str | None = None,
    output_override: str | None = None,
) -> ExperimentConfig:
    input_source = experiment_config.input_source
    if input_path_override is not None:
        input_source = replace(
            input_source,
            path=Path(input_path_override).expanduser().resolve(),
            values=None,
        )
    if input_type_override is not None:
        input_source = replace(input_source, input_type=_normalize_input_type(input_type_override))

    output_config = experiment_config.output
    if output_override is not None:
        output_config = replace(output_config, metrics_file=_resolve_output_target(output_override, "metrics.json"))

    return replace(experiment_config, input_source=input_source, output=output_config)


def apply_trace_config_defaults(trace_config: TraceConfig, experiment_config: ExperimentConfig) -> TraceConfig:
    effective_trace_enabled = trace_config.trace_enabled
    effective_benchmark_enabled = trace_config.benchmark_enabled
    effective_mode = trace_config.mode
    if not trace_config.mode_overridden and not trace_config.trace_enabled and not trace_config.benchmark_enabled:
        default_mode = experiment_config.default_mode()
        effective_mode = default_mode
        if default_mode == "trace":
            effective_trace_enabled = True
        elif default_mode == "benchmark":
            effective_benchmark_enabled = True

    effective_ticks = trace_config.ticks if trace_config.ticks_overridden else max(
        1,
        int(experiment_config.runtime.ticks or trace_config.ticks),
    )
    effective_trace_format = (
        trace_config.trace_format
        if trace_config.trace_format_overridden
        else str(experiment_config.runtime.trace_format or trace_config.trace_format)
    )

    effective_out_target = trace_config.out_target
    if (
        effective_trace_enabled
        and not trace_config.out_target_overridden
        and experiment_config.output.trace_file is not None
    ):
        effective_out_target = str(
            _resolve_output_target(
                experiment_config.output.trace_file,
                "trace.jsonl" if _normalize_trace_format(effective_trace_format) == "jsonl" else "trace.log",
            )
        )
    if trace_config.generic_output_target is not None and not trace_config.out_target_overridden and effective_trace_enabled:
        effective_out_target = str(
            _resolve_output_target(
                trace_config.generic_output_target,
                "trace.jsonl" if _normalize_trace_format(effective_trace_format) == "jsonl" else "trace.log",
            )
        )

    effective_benchmark_output_target = trace_config.benchmark_output_target
    if (
        effective_benchmark_enabled
        and not trace_config.benchmark_output_overridden
        and experiment_config.output.benchmark_file is not None
    ):
        effective_benchmark_output_target = str(
            _resolve_output_target(
                experiment_config.output.benchmark_file,
                "benchmark.jsonl" if trace_config.benchmark_format == "jsonl" else "benchmark.log",
            )
        )
    if trace_config.generic_output_target is not None and not trace_config.benchmark_output_overridden and effective_benchmark_enabled:
        effective_benchmark_output_target = str(
            _resolve_output_target(
                trace_config.generic_output_target,
                "benchmark.jsonl" if trace_config.benchmark_format == "jsonl" else "benchmark.log",
            )
        )

    return replace(
        trace_config,
        trace_enabled=effective_trace_enabled,
        benchmark_enabled=effective_benchmark_enabled,
        mode=effective_mode,
        scenario=experiment_config.name,
        ticks=effective_ticks,
        trace_format=_normalize_trace_format(effective_trace_format),
        out_target=effective_out_target,
        benchmark_output_target=effective_benchmark_output_target,
    )


def _parse_runtime_defaults(runtime_payload: dict[str, Any]) -> RuntimeDefaultsConfig:
    ticks_value = runtime_payload.get("ticks")
    return RuntimeDefaultsConfig(
        ticks=int(ticks_value) if ticks_value is not None else None,
        trace=bool(runtime_payload.get("trace", False)),
        trace_format=_normalize_trace_format(runtime_payload.get("trace_format")),
        benchmark=bool(runtime_payload.get("benchmark", False)),
        interactive=bool(runtime_payload.get("interactive", False)),
    )


def _parse_interactive_defaults(interactive_payload: dict[str, Any]) -> InteractiveDefaultsConfig:
    return InteractiveDefaultsConfig(
        output=str(interactive_payload.get("output", "compact")).strip().lower(),
        trace=bool(interactive_payload.get("trace", False)),
        normalize_inputs=str(interactive_payload.get("normalize_inputs", "normalized")).strip().lower(),
    )


def _parse_input_space(input_space_payload: dict[str, Any] | None) -> InputSpaceConfig:
    payload = dict(input_space_payload or {})
    rank_value = int(payload.get("rank", 1))
    if rank_value < 1 or rank_value > 3:
        raise ValueError(f"Input-space rank must be 1, 2, or 3. Got: {rank_value}")

    port_names = tuple(parse_port_names(payload.get("ports")))
    raw_shape = payload.get("shape")
    if raw_shape is None:
        shape_values = (len(port_names),) if rank_value == 1 else ()
    else:
        if not isinstance(raw_shape, list) or not raw_shape:
            raise ValueError("Input-space shape must be a non-empty list when provided.")
        shape_values = tuple(int(shape_value) for shape_value in raw_shape)

    if len(shape_values) != rank_value:
        raise ValueError(
            f"Input-space rank {rank_value} requires a shape with {rank_value} dimensions; got {list(shape_values)}."
        )
    for dimension_value in shape_values:
        if dimension_value <= 0:
            raise ValueError(f"Input-space dimensions must be positive. Got: {list(shape_values)}")
    if rank_value > 1 and len(port_names) != 1:
        raise ValueError("This PR supports exactly one configured port for rank-2 and rank-3 input spaces.")

    return InputSpaceConfig(
        rank=rank_value,
        shape=shape_values,
        ports=port_names,
        gain=float(payload.get("gain", 1.0)),
        epsilon_fire=float(payload.get("epsilon_fire", 0.01)),
    )


def _parse_regions(network_payload: dict[str, Any], input_ports: list[str]) -> list[RegionScaffold]:
    regions_payload = network_payload.get("regions")
    if not regions_payload:
        raise ValueError("Experiment config must define network.regions.")
    return [parse_region_scaffold(region_payload, input_ports) for region_payload in regions_payload]


def _parse_input_source(input_payload: dict[str, Any], base_directory: Path) -> InputSourceConfig:
    raw_resize = input_payload.get("resize")
    if raw_resize is not None:
        if not isinstance(raw_resize, list) or not raw_resize:
            raise ValueError("Image resize must be a non-empty list when provided.")
        resize = tuple(int(dimension_value) for dimension_value in raw_resize)
    else:
        resize = None
    return InputSourceConfig(
        input_type=_normalize_input_type(input_payload.get("type")),
        path=_resolve_optional_path(base_directory, input_payload.get("path")),
        values=input_payload.get("values"),
        mode=str(input_payload.get("mode", "grayscale")).strip().lower(),
        resize=resize,
        sample_rate=int(input_payload["sample_rate"]) if input_payload.get("sample_rate") is not None else None,
        frame_size=int(input_payload["frame_size"]) if input_payload.get("frame_size") is not None else None,
        hop_size=int(input_payload["hop_size"]) if input_payload.get("hop_size") is not None else None,
    )


def _parse_output_config(output_payload: dict[str, Any], base_directory: Path) -> OutputConfig:
    return OutputConfig(
        trace_file=_resolve_optional_path(base_directory, output_payload.get("trace_file")),
        metrics_file=_resolve_optional_path(base_directory, output_payload.get("metrics_file")),
        benchmark_file=_resolve_optional_path(
            base_directory,
            output_payload.get("benchmark_output_file") or output_payload.get("benchmark_file"),
        ),
    )


def _normalize_input_type(raw_input_type: Any) -> str:
    normalized_input_type = str(raw_input_type or "sequence").strip().lower()
    if normalized_input_type not in {"sequence", "tensor", "image", "audio"}:
        raise ValueError(f"Unsupported input type: {raw_input_type}")
    return normalized_input_type


def _normalize_trace_format(raw_trace_format: Any) -> str:
    return "jsonl" if str(raw_trace_format or "text").strip().lower() == "jsonl" else "text"


def _resolve_optional_path(base_directory: Path, raw_path: Any) -> Path | None:
    if raw_path is None or not str(raw_path).strip():
        return None
    candidate_path = Path(str(raw_path)).expanduser()
    if not candidate_path.is_absolute():
        candidate_path = (base_directory / candidate_path).resolve()
    else:
        candidate_path = candidate_path.resolve()
    return candidate_path


def _resolve_output_target(raw_target: Any, default_filename: str) -> Path:
    candidate_path = Path(str(raw_target)).expanduser()
    if not candidate_path.is_absolute():
        candidate_path = candidate_path.resolve()
    else:
        candidate_path = candidate_path.resolve()

    if (candidate_path.exists() and candidate_path.is_dir()) or candidate_path.suffix == "":
        return candidate_path / default_filename
    return candidate_path
