from __future__ import annotations

import csv
import json
import math
import struct
import wave
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .experiment_config import ExperimentConfig, InputSpaceConfig


@dataclass(frozen=True)
class LoadedInputSchedule:
    input_type: str
    rank: int
    tick_payloads: list[Any]
    source_path: Path | None = None


def load_input_schedule(experiment_config: ExperimentConfig) -> LoadedInputSchedule:
    input_source = experiment_config.input_source
    if input_source.input_type == "sequence":
        tick_payloads = _load_sequence_payloads(experiment_config)
    elif input_source.input_type == "tensor":
        tick_payloads = _load_tensor_payloads(experiment_config)
    elif input_source.input_type == "image":
        tick_payloads = _load_image_payloads(experiment_config)
    elif input_source.input_type == "audio":
        tick_payloads = _load_audio_payloads(experiment_config)
    else:
        raise ValueError(f"Unsupported input type: {input_source.input_type}")
    return LoadedInputSchedule(
        input_type=input_source.input_type,
        rank=experiment_config.input_space.rank,
        tick_payloads=tick_payloads,
        source_path=input_source.path,
    )


def expand_tick_payloads(
    experiment_config: ExperimentConfig,
    loaded_schedule: LoadedInputSchedule,
    requested_ticks: int | None,
) -> list[Any]:
    if not loaded_schedule.tick_payloads:
        raise ValueError("Configured input source did not produce any tick payloads.")

    if requested_ticks is None or requested_ticks <= 0:
        return list(loaded_schedule.tick_payloads)
    if requested_ticks <= len(loaded_schedule.tick_payloads):
        return list(loaded_schedule.tick_payloads[:requested_ticks])

    if len(loaded_schedule.tick_payloads) == 1 and loaded_schedule.input_type in {"tensor", "image"}:
        return [loaded_schedule.tick_payloads[0] for _repeat_index in range(requested_ticks)]

    raise ValueError(
        "Requested ticks exceed available input payloads. "
        f"requested_ticks={requested_ticks} available_ticks={len(loaded_schedule.tick_payloads)}"
    )


def flatten_tensor_payload(tensor_payload: Any) -> tuple[list[float], list[int]]:
    inferred_shape = infer_tensor_shape(tensor_payload)
    flat_values: list[float] = []
    _append_flat_values(tensor_payload, flat_values)
    return flat_values, inferred_shape


def infer_tensor_shape(tensor_payload: Any) -> list[int]:
    if not isinstance(tensor_payload, list):
        return []
    if not tensor_payload:
        return [0]
    first_shape = infer_tensor_shape(tensor_payload[0])
    for child_payload in tensor_payload[1:]:
        child_shape = infer_tensor_shape(child_payload)
        if child_shape != first_shape:
            raise ValueError("Tensor payload must be rectangular.")
    return [len(tensor_payload)] + first_shape


def _load_sequence_payloads(experiment_config: ExperimentConfig) -> list[Any]:
    raw_payload = _load_values_or_path(experiment_config.input_source.values, experiment_config.input_source.path)
    input_space = experiment_config.input_space
    if input_space.uses_scalar_port_inputs():
        return _normalize_scalar_sequence(raw_payload, input_space)
    return _normalize_tensor_sequence(raw_payload, input_space)


def _load_tensor_payloads(experiment_config: ExperimentConfig) -> list[Any]:
    raw_payload = _load_values_or_path(experiment_config.input_source.values, experiment_config.input_source.path)
    return _normalize_tensor_ticks(raw_payload, experiment_config.input_space)


def _load_image_payloads(experiment_config: ExperimentConfig) -> list[Any]:
    input_source = experiment_config.input_source
    if input_source.path is None:
        raise ValueError("Image input requires input.path.")
    if experiment_config.input_space.rank != 2:
        raise ValueError("Image input currently requires input_space.rank = 2.")
    image_matrix = load_image_matrix(input_source.path, input_source.mode)
    if input_source.resize is not None:
        if len(input_source.resize) != 2:
            raise ValueError("Image resize must contain [height, width].")
        image_matrix = resize_matrix(image_matrix, int(input_source.resize[0]), int(input_source.resize[1]))
    _validate_matrix_shape(image_matrix, list(experiment_config.input_space.shape))
    return [image_matrix]


def _load_audio_payloads(experiment_config: ExperimentConfig) -> list[Any]:
    input_source = experiment_config.input_source
    if input_source.path is None:
        raise ValueError("Audio input requires input.path.")
    if experiment_config.input_space.rank != 1:
        raise ValueError("Audio input currently requires input_space.rank = 1.")
    frame_size = int(input_source.frame_size or 0)
    hop_size = int(input_source.hop_size or 0)
    if frame_size <= 0 or hop_size <= 0:
        raise ValueError("Audio input requires positive frame_size and hop_size.")
    expected_shape = list(experiment_config.input_space.shape)
    if expected_shape != [frame_size]:
        raise ValueError(
            f"Audio input_space.shape must match frame_size. expected={[frame_size]} configured={expected_shape}"
        )
    return load_audio_windows(input_source.path, input_source.sample_rate, frame_size, hop_size)


def _load_values_or_path(inline_values: Any, input_path: Path | None) -> Any:
    if inline_values is not None:
        return inline_values
    if input_path is None:
        raise ValueError("Input configuration requires either input.values or input.path.")
    return _load_path_payload(input_path)


def _load_path_payload(input_path: Path) -> Any:
    if not input_path.exists():
        raise ValueError(f"Input file not found: {input_path}")

    suffix_value = input_path.suffix.lower()
    if suffix_value == ".json":
        with input_path.open("r", encoding="utf-8") as input_handle:
            return json.load(input_handle)
    if suffix_value == ".jsonl":
        payloads: list[Any] = []
        with input_path.open("r", encoding="utf-8") as input_handle:
            for raw_line in input_handle:
                stripped_line = raw_line.strip()
                if stripped_line:
                    payloads.append(json.loads(stripped_line))
        return payloads
    if suffix_value in {".csv", ".txt"}:
        with input_path.open("r", encoding="utf-8") as input_handle:
            raw_text = input_handle.read().strip()
        if not raw_text:
            return []
        if suffix_value == ".csv":
            csv_reader = csv.reader(raw_text.splitlines())
            flattened_values: list[float] = []
            for row_values in csv_reader:
                for cell_value in row_values:
                    stripped_value = str(cell_value).strip()
                    if stripped_value:
                        flattened_values.append(float(stripped_value))
            return flattened_values
        return [
            float(token_value.strip())
            for token_value in raw_text.replace(",", "\n").splitlines()
            if token_value.strip()
        ]
    raise ValueError(f"Unsupported input file type: {input_path}")


def _normalize_scalar_sequence(raw_payload: Any, input_space: InputSpaceConfig) -> list[dict[str, float]]:
    port_names = list(input_space.ports)
    if isinstance(raw_payload, dict):
        return [_coerce_port_mapping(raw_payload, port_names)]

    if not isinstance(raw_payload, list):
        if len(port_names) != 1:
            raise ValueError("Multi-port scalar input requires a list of per-tick values.")
        return [{port_names[0]: float(raw_payload)}]

    if len(port_names) == 1 and all(not isinstance(item_value, (list, tuple, dict)) for item_value in raw_payload):
        return [{port_names[0]: float(item_value)} for item_value in raw_payload]

    tick_payloads: list[dict[str, float]] = []
    for tick_payload in raw_payload:
        tick_payloads.append(_coerce_port_mapping(tick_payload, port_names))
    return tick_payloads


def _normalize_tensor_sequence(raw_payload: Any, input_space: InputSpaceConfig) -> list[Any]:
    expected_tensor_size = input_space.total_size()
    if expected_tensor_size <= 0:
        raise ValueError("Tensor-like sequence input requires a positive input-space shape.")

    if not isinstance(raw_payload, list):
        raise ValueError("Structured sequence input must be a list of per-tick tensors.")

    if raw_payload and all(not isinstance(item_value, list) for item_value in raw_payload) and expected_tensor_size > 1:
        single_tensor = _coerce_tensor_payload(raw_payload, input_space)
        return [single_tensor]

    return [_coerce_tensor_payload(tick_payload, input_space) for tick_payload in raw_payload]


def _normalize_tensor_ticks(raw_payload: Any, input_space: InputSpaceConfig) -> list[Any]:
    if input_space.uses_scalar_port_inputs():
        return _normalize_scalar_tensor_ticks(raw_payload, input_space)

    payload_depth = _nesting_depth(raw_payload)
    if payload_depth == input_space.rank:
        return [_coerce_tensor_payload(raw_payload, input_space)]
    if payload_depth == input_space.rank + 1:
        if not isinstance(raw_payload, list):
            raise ValueError("Multi-tick tensor input must be a list.")
        return [_coerce_tensor_payload(tick_payload, input_space) for tick_payload in raw_payload]
    raise ValueError(
        f"Tensor payload depth {payload_depth} does not match input-space rank {input_space.rank}."
    )


def _normalize_scalar_tensor_ticks(raw_payload: Any, input_space: InputSpaceConfig) -> list[dict[str, float]]:
    port_names = list(input_space.ports)
    payload_depth = _nesting_depth(raw_payload)
    if payload_depth == 1:
        return [_coerce_port_mapping(raw_payload, port_names)]
    if payload_depth == 2 and isinstance(raw_payload, list):
        return [_coerce_port_mapping(tick_payload, port_names) for tick_payload in raw_payload]
    raise ValueError(
        "Tensor input for scalar-port spaces must be a vector or a list of vectors matching the configured ports."
    )


def _coerce_port_mapping(tick_payload: Any, port_names: list[str]) -> dict[str, float]:
    if isinstance(tick_payload, dict):
        mapped_values: dict[str, float] = {}
        for port_name in port_names:
            if port_name not in tick_payload:
                raise ValueError(f"Per-tick input is missing port '{port_name}'.")
            mapped_values[port_name] = float(tick_payload[port_name])
        return mapped_values

    if isinstance(tick_payload, (list, tuple)):
        if len(tick_payload) != len(port_names):
            raise ValueError(
                f"Per-tick input length {len(tick_payload)} does not match configured ports {port_names}."
            )
        return {
            port_names[port_index]: float(tick_payload[port_index])
            for port_index in range(len(port_names))
        }

    if len(port_names) != 1:
        raise ValueError("Multi-port input requires a list/tuple or object mapping ports to values.")
    return {port_names[0]: float(tick_payload)}


def _coerce_tensor_payload(tensor_payload: Any, input_space: InputSpaceConfig) -> Any:
    inferred_shape = infer_tensor_shape(tensor_payload)
    expected_shape = list(input_space.shape)
    if inferred_shape != expected_shape:
        raise ValueError(f"Tensor shape mismatch. expected={expected_shape} received={inferred_shape}")
    return _convert_nested_floats(tensor_payload)


def load_image_matrix(image_path: Path, image_mode: str) -> list[list[float]]:
    suffix_value = image_path.suffix.lower()
    if suffix_value == ".pgm":
        return _load_netpbm_grayscale(image_path)

    try:
        from PIL import Image
    except ImportError as import_error:
        raise ValueError(
            f"Image input currently requires Pillow for {image_path.suffix or 'that format'} files."
        ) from import_error

    with Image.open(image_path) as image_handle:
        if image_mode != "grayscale":
            raise ValueError(f"Unsupported image mode: {image_mode}. Only grayscale is supported in this PR.")
        grayscale_image = image_handle.convert("L")
        image_rows: list[list[float]] = []
        for row_index in range(grayscale_image.height):
            row_values: list[float] = []
            for col_index in range(grayscale_image.width):
                row_values.append(float(grayscale_image.getpixel((col_index, row_index))) / 255.0)
            image_rows.append(row_values)
        return image_rows


def resize_matrix(matrix_values: list[list[float]], target_height: int, target_width: int) -> list[list[float]]:
    if target_height <= 0 or target_width <= 0:
        raise ValueError("Image resize target dimensions must be positive.")
    source_height = len(matrix_values)
    source_width = len(matrix_values[0]) if source_height > 0 else 0
    if source_height == 0 or source_width == 0:
        raise ValueError("Cannot resize an empty image matrix.")

    resized_rows: list[list[float]] = []
    for row_index in range(target_height):
        source_row = min(source_height - 1, int(math.floor(row_index * source_height / target_height)))
        row_values: list[float] = []
        for col_index in range(target_width):
            source_col = min(source_width - 1, int(math.floor(col_index * source_width / target_width)))
            row_values.append(float(matrix_values[source_row][source_col]))
        resized_rows.append(row_values)
    return resized_rows


def load_audio_windows(
    audio_path: Path,
    expected_sample_rate: int | None,
    frame_size: int,
    hop_size: int,
) -> list[list[float]]:
    with wave.open(str(audio_path), "rb") as wave_handle:
        channel_count = int(wave_handle.getnchannels())
        sample_width = int(wave_handle.getsampwidth())
        sample_rate = int(wave_handle.getframerate())
        frame_count = int(wave_handle.getnframes())
        raw_frames = wave_handle.readframes(frame_count)

    if expected_sample_rate is not None and sample_rate != int(expected_sample_rate):
        raise ValueError(
            f"Audio sample rate {sample_rate} does not match configured sample_rate {expected_sample_rate}. "
            "Resampling is not implemented in this PR."
        )

    samples = _decode_pcm_samples(raw_frames, sample_width, channel_count)
    if not samples:
        return [[0.0 for _sample_index in range(frame_size)]]

    windowed_samples: list[list[float]] = []
    start_index = 0
    last_start_index = max(0, len(samples) - frame_size)
    while start_index <= last_start_index:
        frame_values = samples[start_index:start_index + frame_size]
        if len(frame_values) < frame_size:
            frame_values = frame_values + [0.0 for _pad_index in range(frame_size - len(frame_values))]
        windowed_samples.append(frame_values)
        start_index += hop_size

    if not windowed_samples:
        padded_frame = samples[:frame_size]
        if len(padded_frame) < frame_size:
            padded_frame = padded_frame + [0.0 for _pad_index in range(frame_size - len(padded_frame))]
        windowed_samples.append(padded_frame)
    return windowed_samples


def _decode_pcm_samples(raw_frames: bytes, sample_width: int, channel_count: int) -> list[float]:
    if sample_width == 1:
        raw_values = [float(frame_byte - 128) / 128.0 for frame_byte in raw_frames]
    elif sample_width == 2:
        sample_count = len(raw_frames) // 2
        raw_values = [float(sample_value) / 32768.0 for sample_value in struct.unpack(f"<{sample_count}h", raw_frames)]
    elif sample_width == 4:
        sample_count = len(raw_frames) // 4
        raw_values = [float(sample_value) / 2147483648.0 for sample_value in struct.unpack(f"<{sample_count}i", raw_frames)]
    else:
        raise ValueError(f"Unsupported audio sample width: {sample_width}")

    if channel_count <= 1:
        return raw_values

    mono_values: list[float] = []
    for frame_offset in range(0, len(raw_values), channel_count):
        channel_window = raw_values[frame_offset:frame_offset + channel_count]
        mono_values.append(sum(channel_window) / len(channel_window))
    return mono_values


def _load_netpbm_grayscale(image_path: Path) -> list[list[float]]:
    with image_path.open("rb") as image_handle:
        raw_content = image_handle.read()

    tokens: list[bytes] = []
    token_buffer = b""
    comment_mode = False
    for byte_value in raw_content:
        if comment_mode:
            if byte_value in {10, 13}:
                comment_mode = False
            continue
        if byte_value == 35:
            comment_mode = True
            continue
        if byte_value in {9, 10, 13, 32}:
            if token_buffer:
                tokens.append(token_buffer)
                token_buffer = b""
            continue
        token_buffer += bytes([byte_value])
    if token_buffer:
        tokens.append(token_buffer)

    if len(tokens) < 4:
        raise ValueError(f"Invalid NetPBM image file: {image_path}")
    magic_value = tokens[0].decode("ascii")
    if magic_value not in {"P2", "P5"}:
        raise ValueError(f"Unsupported NetPBM image format: {magic_value}")
    width_value = int(tokens[1])
    height_value = int(tokens[2])
    max_value = int(tokens[3])
    if max_value <= 0:
        raise ValueError(f"Invalid NetPBM max value: {max_value}")

    pixel_values: list[float] = []
    if magic_value == "P2":
        for token_value in tokens[4:]:
            pixel_values.append(float(int(token_value)) / float(max_value))
    else:
        header_token_count = 4
        header_occurrences = 0
        data_start_index = 0
        current_comment_mode = False
        current_token = b""
        for byte_index, byte_value in enumerate(raw_content):
            if current_comment_mode:
                if byte_value in {10, 13}:
                    current_comment_mode = False
                continue
            if byte_value == 35:
                current_comment_mode = True
                continue
            if byte_value in {9, 10, 13, 32}:
                if current_token:
                    header_occurrences += 1
                    current_token = b""
                    if header_occurrences == header_token_count:
                        data_start_index = byte_index + 1
                        break
                continue
            current_token += bytes([byte_value])
        raw_pixels = raw_content[data_start_index:]
        pixel_values = [float(pixel_byte) / float(max_value) for pixel_byte in raw_pixels[: width_value * height_value]]

    if len(pixel_values) != width_value * height_value:
        raise ValueError(f"Unexpected pixel count in NetPBM image: {image_path}")

    image_rows: list[list[float]] = []
    for row_index in range(height_value):
        row_start = row_index * width_value
        image_rows.append(pixel_values[row_start:row_start + width_value])
    return image_rows


def _validate_matrix_shape(matrix_values: list[list[float]], expected_shape: list[int]) -> None:
    actual_shape = infer_tensor_shape(matrix_values)
    if actual_shape != expected_shape:
        raise ValueError(f"Image shape mismatch. expected={expected_shape} received={actual_shape}")


def _append_flat_values(tensor_payload: Any, flat_values: list[float]) -> None:
    if isinstance(tensor_payload, list):
        for child_payload in tensor_payload:
            _append_flat_values(child_payload, flat_values)
        return
    flat_values.append(float(tensor_payload))


def _convert_nested_floats(tensor_payload: Any) -> Any:
    if isinstance(tensor_payload, list):
        return [_convert_nested_floats(child_payload) for child_payload in tensor_payload]
    return float(tensor_payload)


def _nesting_depth(raw_payload: Any) -> int:
    if not isinstance(raw_payload, list):
        return 0
    if not raw_payload:
        return 1
    return 1 + _nesting_depth(raw_payload[0])
