from __future__ import annotations

from pathlib import Path

import numpy as np

from .fitting import evaluate_line
from .io import write_colored_ply
from .reference import evaluate_reference_line
from .structures import FittedLine, ReferenceLine, TrackSolution


COLOR_BASE = np.array([170, 170, 170], dtype=np.uint8)
COLOR_LEFT_BAND = np.array([0, 114, 189], dtype=np.uint8)
COLOR_RIGHT_BAND = np.array([217, 83, 25], dtype=np.uint8)
COLOR_LEFT_REP_LINE = np.array([0, 220, 255], dtype=np.uint8)
COLOR_RIGHT_REP_LINE = np.array([255, 196, 0], dtype=np.uint8)
COLOR_LEFT_REF_LINE = np.array([0, 200, 120], dtype=np.uint8)
COLOR_RIGHT_REF_LINE = np.array([200, 0, 160], dtype=np.uint8)
COLOR_LEFT_MEASURE = np.array([0, 255, 120], dtype=np.uint8)
COLOR_RIGHT_MEASURE = np.array([255, 80, 200], dtype=np.uint8)
COLOR_CORRESPONDENCE = np.array([255, 255, 255], dtype=np.uint8)
COLOR_ANCHOR = np.array([255, 64, 64], dtype=np.uint8)


def _points_and_colors(
    points: np.ndarray,
    color: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    if len(points) == 0:
        return np.zeros((0, 3), dtype=np.float64), np.zeros((0, 3), dtype=np.uint8)
    colors = np.repeat(color[None, :], len(points), axis=0)
    return points, colors


def _sample_segment(
    start: np.ndarray,
    end: np.ndarray,
    step_mm: float,
) -> np.ndarray:
    start = np.asarray(start, dtype=np.float64)
    end = np.asarray(end, dtype=np.float64)
    length = float(np.linalg.norm(end - start))
    if length <= 1e-9:
        return start[None, :]
    count = max(int(np.ceil(length / max(step_mm, 1.0))) + 1, 2)
    t = np.linspace(0.0, 1.0, count)
    return start[None, :] + (end - start)[None, :] * t[:, None]


def _sample_polyline(points: np.ndarray, step_mm: float) -> np.ndarray:
    if len(points) == 0:
        return np.zeros((0, 3), dtype=np.float64)
    if len(points) == 1:
        return points.copy()
    segments = []
    for start, end in zip(points[:-1], points[1:], strict=True):
        segment = _sample_segment(start, end, step_mm)
        if segments:
            segment = segment[1:]
        segments.append(segment)
    return np.vstack(segments)


def _marker_cross(
    center: np.ndarray,
    size_mm: float,
    step_mm: float,
) -> np.ndarray:
    center = np.asarray(center, dtype=np.float64)
    offsets = np.arange(-size_mm, size_mm + step_mm, step_mm)
    samples = [center[None, :]]
    for axis in range(3):
        points = np.repeat(center[None, :], len(offsets), axis=0)
        points[:, axis] += offsets
        samples.append(points)
    return np.vstack(samples)


def _track_band_mask(
    points: np.ndarray,
    track: TrackSolution | None,
    config: dict,
) -> np.ndarray:
    if track is None:
        return np.zeros(len(points), dtype=bool)
    rep_cfg = config["representative"]
    mask = np.zeros(len(points), dtype=bool)
    for candidate in track.candidates:
        candidate_mask = (
            (points[:, 0] >= candidate.x_start)
            & (points[:, 0] < candidate.x_end)
            & (points[:, 1] >= candidate.y_min - 0.5 * rep_cfg["band_y_margin_mm"])
            & (points[:, 1] <= candidate.y_max + 0.5 * rep_cfg["band_y_margin_mm"])
            & (points[:, 2] <= candidate.z_low + rep_cfg["band_z_margin_mm"])
        )
        mask |= candidate_mask
    return mask


def _sample_fitted_line(
    line: FittedLine | None,
    step_mm: float,
) -> np.ndarray:
    if line is None:
        return np.zeros((0, 3), dtype=np.float64)
    x_min = float(line.x_support.min())
    x_max = float(line.x_support.max())
    x_grid = np.arange(x_min, x_max + step_mm, step_mm)
    y, z, conf = evaluate_line(line, x_grid)
    keep = conf > 0.0
    return np.column_stack([x_grid[keep], y[keep], z[keep]])


def _sample_reference_line(
    reference_line: ReferenceLine,
    step_mm: float,
) -> np.ndarray:
    if not reference_line.available or reference_line.point_front is None or reference_line.point_back is None:
        return np.zeros((0, 3), dtype=np.float64)
    return _sample_segment(reference_line.point_front, reference_line.point_back, step_mm)


def _measurement_overlay_geometry(
    measurement: dict | None,
    measure_color: np.ndarray,
    step_mm: float,
    marker_size_mm: float,
) -> list[tuple[np.ndarray, np.ndarray]]:
    layers: list[tuple[np.ndarray, np.ndarray]] = []
    if not measurement or not measurement.get("available", False):
        return layers
    beam_point = np.asarray(measurement["beam_point"], dtype=np.float64)
    ref_point = np.asarray(measurement["reference_point"], dtype=np.float64)
    vertical_anchor = np.asarray(measurement["vertical_anchor_point"], dtype=np.float64)

    layers.append(_points_and_colors(_marker_cross(beam_point, marker_size_mm, step_mm), measure_color))
    layers.append(_points_and_colors(_marker_cross(ref_point, marker_size_mm, step_mm), COLOR_CORRESPONDENCE))
    layers.append(_points_and_colors(_marker_cross(vertical_anchor, marker_size_mm, step_mm), COLOR_ANCHOR))
    layers.append(_points_and_colors(_sample_segment(beam_point, vertical_anchor, step_mm), measure_color))
    layers.append(_points_and_colors(_sample_segment(vertical_anchor, ref_point, step_mm), COLOR_CORRESPONDENCE))
    return layers


def export_measurement_overlay_plys(
    output_dir: str | Path,
    roi_points: np.ndarray,
    left_track: TrackSolution | None,
    right_track: TrackSolution | None,
    left_line: FittedLine | None,
    right_line: FittedLine | None,
    left_reference: ReferenceLine,
    right_reference: ReferenceLine,
    left_measurement: dict | None,
    right_measurement: dict | None,
    config: dict,
    coordinate_transform=None,
) -> dict[str, str]:
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    step_mm = max(config["runtime"]["spacing_mm"] * 0.8, 20.0)
    marker_size_mm = max(config["runtime"]["spacing_mm"] * 2.5, 60.0)

    roi_colors = np.repeat(COLOR_BASE[None, :], len(roi_points), axis=0)
    roi_path = output_dir / "roi_cropped_cloud.ply"

    left_band_mask = _track_band_mask(roi_points, left_track, config)
    right_band_mask = _track_band_mask(roi_points, right_track, config)
    left_band_points = roi_points[left_band_mask]
    right_band_points = roi_points[right_band_mask]

    left_rep_line = _sample_fitted_line(left_line, step_mm)
    right_rep_line = _sample_fitted_line(right_line, step_mm)
    left_ref_line = _sample_reference_line(left_reference, step_mm)
    right_ref_line = _sample_reference_line(right_reference, step_mm)

    def build_overlay(include_left: bool, include_right: bool) -> tuple[np.ndarray, np.ndarray]:
        layers: list[tuple[np.ndarray, np.ndarray]] = [(_points_and_colors(roi_points, COLOR_BASE))]
        if include_left:
            layers.append(_points_and_colors(left_band_points, COLOR_LEFT_BAND))
            layers.append(_points_and_colors(left_rep_line, COLOR_LEFT_REP_LINE))
            layers.append(_points_and_colors(left_ref_line, COLOR_LEFT_REF_LINE))
            layers.extend(_measurement_overlay_geometry(left_measurement, COLOR_LEFT_MEASURE, step_mm, marker_size_mm))
        if include_right:
            layers.append(_points_and_colors(right_band_points, COLOR_RIGHT_BAND))
            layers.append(_points_and_colors(right_rep_line, COLOR_RIGHT_REP_LINE))
            layers.append(_points_and_colors(right_ref_line, COLOR_RIGHT_REF_LINE))
            layers.extend(_measurement_overlay_geometry(right_measurement, COLOR_RIGHT_MEASURE, step_mm, marker_size_mm))
        points = np.vstack([layer_points for layer_points, _ in layers if len(layer_points) > 0])
        colors = np.vstack([layer_colors for layer_points, layer_colors in layers if len(layer_points) > 0])
        return points, colors

    full_points, full_colors = build_overlay(include_left=True, include_right=True)
    left_points, left_colors = build_overlay(include_left=True, include_right=False)
    right_points, right_colors = build_overlay(include_left=False, include_right=True)
    if coordinate_transform is not None:
        roi_points_out = coordinate_transform(roi_points)
        full_points = coordinate_transform(full_points)
        left_points = coordinate_transform(left_points)
        right_points = coordinate_transform(right_points)
    else:
        roi_points_out = roi_points

    full_path = output_dir / "roi_measurement_overlay_full.ply"
    left_path = output_dir / "roi_measurement_overlay_left.ply"
    right_path = output_dir / "roi_measurement_overlay_right.ply"
    write_colored_ply(roi_path, roi_points_out, roi_colors)
    write_colored_ply(full_path, full_points, full_colors)
    write_colored_ply(left_path, left_points, left_colors)
    write_colored_ply(right_path, right_points, right_colors)
    return {
        "roi_cropped_cloud": str(roi_path),
        "roi_measurement_overlay_full": str(full_path),
        "roi_measurement_overlay_left": str(left_path),
        "roi_measurement_overlay_right": str(right_path),
    }