"""Read YOLO-seg .txt predictions and build a per-tuber morphometrics dataframe.

YOLO-seg .txt format (one detection per line):
    class_id  x1 y1  x2 y2  ...  xN yN     # coords normalized to [0, 1]
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image

from .morphometrics import polygon_shape_metrics

IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"}


def read_yolo_seg_txt(txt_path: str | Path) -> list[dict]:
    """Parse a YOLO-seg .txt file into a list of {class_id, xy_norm}.

    Lines with fewer than 3 vertices (6 coordinates) or an odd coordinate
    count are skipped.
    """
    polys: list[dict] = []
    with open(txt_path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            class_id = int(float(parts[0]))
            coords = np.array(list(map(float, parts[1:])), dtype=float)
            if coords.size < 6 or coords.size % 2 != 0:
                continue
            polys.append({"class_id": class_id, "xy_norm": coords.reshape(-1, 2)})
    return polys


def _list_images(folder: Path) -> list[Path]:
    return sorted(p for p in folder.rglob("*") if p.suffix.lower() in IMG_EXTS)


def build_morphometrics_dataframe(
    trial_name: str,
    images_dir: str | Path,
    preds_dir: str | Path,
) -> pd.DataFrame:
    """Produce one row per detected tuber with all pixel-scale shape metrics.

    Pairs each .txt file in `preds_dir` with the image of the same stem in
    `images_dir` so we can read the image dimensions (needed both for
    denormalizing polygon coordinates and for the downstream calibration).
    """
    images_dir = Path(images_dir)
    preds_dir = Path(preds_dir)

    if not preds_dir.exists():
        raise FileNotFoundError(f"[{trial_name}] preds_dir not found: {preds_dir}")
    if not images_dir.exists():
        raise FileNotFoundError(f"[{trial_name}] images_dir not found: {images_dir}")

    image_map = {p.stem: p for p in _list_images(images_dir)}
    txt_files = sorted(preds_dir.glob("*.txt"))
    if not txt_files:
        raise ValueError(f"[{trial_name}] No .txt files found in {preds_dir}")

    rows: list[dict] = []
    skipped_txt_no_image = 0

    for txt in txt_files:
        img_stem = txt.stem
        img_path = image_map.get(img_stem)
        if img_path is None:
            skipped_txt_no_image += 1
            continue

        with Image.open(img_path) as im:
            img_width, img_height = im.size

        for line_idx, poly in enumerate(read_yolo_seg_txt(txt), start=1):
            xy = poly["xy_norm"]
            xs = xy[:, 0] * img_width
            ys = xy[:, 1] * img_height
            metrics = polygon_shape_metrics(xs, ys)

            rows.append(
                {
                    "trial": trial_name,
                    "image_id": img_stem,
                    "detection_id": f"{img_stem}_{line_idx}",
                    "class_id": poly["class_id"],
                    "img_width_px": img_width,
                    "img_height_px": img_height,
                    "n_points": len(xs),
                    **metrics,
                }
            )

    df = pd.DataFrame(rows)
    if skipped_txt_no_image:
        print(
            f"[{trial_name}] images={len(image_map)} txt={len(txt_files)} "
            f"instances={len(df)} skipped_txt_no_image={skipped_txt_no_image}"
        )
    return df


def plot_id_from_image_id(image_id: pd.Series) -> pd.Series:
    """Derive a raw plot identifier from the last 8 characters of the image stem.

    The result (e.g. "NCPT-129", "5_ADC101") is further normalized by
    standardize_plot_id in the yield-prediction stage.
    """
    return image_id.str[-8:]