# filename: src/data/preprocessor.py
# purpose:  YOLO annotation validation, data.yaml generation, dataset structure verification
# version:  1.0

import logging
from pathlib import Path

import yaml

from config import CLASSES, NUM_CLASSES, DATA_PROCESSED_DIR

logger = logging.getLogger(__name__)


def validate_yolo_annotations(labels_dir: Path) -> None:
    """
    Assert all YOLO label files have valid normalized values.
    Raises AssertionError immediately on the first invalid line.
    Call this before every YOLO training run.
    """
    labels_dir = Path(labels_dir)
    if not labels_dir.exists():
        logger.warning(f"Labels dir not found: {labels_dir}")
        return

    files_checked = 0
    lines_checked = 0

    for txt_file in labels_dir.glob("*.txt"):
        content = txt_file.read_text(encoding="utf-8").strip()
        if not content:
            continue
        for line in content.splitlines():
            parts = line.split()
            if len(parts) < 5:
                continue
            class_id = int(parts[0])
            x_c = float(parts[1])
            y_c = float(parts[2])
            w = float(parts[3])
            h = float(parts[4])

            assert 0 <= class_id < NUM_CLASSES, (
                f"class_id {class_id} out of range [0, {NUM_CLASSES}) in {txt_file}"
            )
            # Use <= 1.0 — bbox_to_yolo clamps to min(1.0,...) so boundary values are valid
            assert 0.0 < x_c <= 1.0, f"x_center {x_c} not in (0,1] in {txt_file}"
            assert 0.0 < y_c <= 1.0, f"y_center {y_c} not in (0,1] in {txt_file}"
            assert 0.0 < w <= 1.0, f"width {w} not in (0,1] in {txt_file}"
            assert 0.0 < h <= 1.0, f"height {h} not in (0,1] in {txt_file}"
            # Catch clearly wrong values (negative or >1)
            assert x_c >= 0 and y_c >= 0, f"Negative center in {txt_file}"
            lines_checked += 1
        files_checked += 1

    print(
        f"YOLO annotations valid: {files_checked} files, {lines_checked} annotations in {labels_dir.name}/"
    )


def create_yolo_data_yaml(detection_dir: Path) -> Path:
    """
    Generate data.yaml with absolute paths — works identically in terminal and Colab.
    Always regenerate before YOLO training (Colab paths differ from local paths).
    """
    detection_dir = Path(detection_dir)
    config = {
        "path": str(detection_dir.absolute()),
        "train": "images/train",
        "val": "images/val",
        "nc": NUM_CLASSES,
        "names": {i: cls for i, cls in enumerate(CLASSES)},
    }
    yaml_path = detection_dir / "data.yaml"
    with open(yaml_path, "w", encoding="utf-8") as f:
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)

    # Verify it loads cleanly
    with open(yaml_path, encoding="utf-8") as f:
        loaded = yaml.safe_load(f)
    assert loaded["nc"] == NUM_CLASSES
    assert len(loaded["names"]) == NUM_CLASSES

    logger.info(f"data.yaml written: {yaml_path}")
    print(f"data.yaml OK: nc={loaded['nc']}, path={loaded['path']}")
    return yaml_path


def verify_dataset_structure() -> dict:
    """
    Count images in every split/class folder and return a summary dict.
    Confirms the dataset was built correctly before preprocessing.
    """

    cls_dir = DATA_PROCESSED_DIR / "classification"
    det_dir = DATA_PROCESSED_DIR / "detection"

    result: dict = {
        "classification": {},
        "detection": {},
        "issues": [],
    }

    # Classification counts
    for split in ["train", "val", "test"]:
        split_total = 0
        per_class: dict[str, int] = {}
        for cls in CLASSES:
            count = (
                len(list((cls_dir / split / cls).glob("*.jpg")))
                if (cls_dir / split / cls).exists()
                else 0
            )
            per_class[cls] = count
            split_total += count
            if count == 0:
                result["issues"].append(f"classification/{split}/{cls}: 0 images")
        result["classification"][split] = {"total": split_total, "per_class": per_class}

    # Detection counts
    for split in ["train", "val"]:
        imgs = (
            len(list((det_dir / "images" / split).glob("*.jpg")))
            if (det_dir / "images" / split).exists()
            else 0
        )
        lbls = (
            len(list((det_dir / "labels" / split).glob("*.txt")))
            if (det_dir / "labels" / split).exists()
            else 0
        )
        result["detection"][split] = {"images": imgs, "labels": lbls}
        if imgs != lbls:
            result["issues"].append(
                f"detection/{split}: images({imgs}) != labels({lbls})"
            )

    # Expected totals
    total_cls = sum(
        result["classification"][s]["total"] for s in ["train", "val", "test"]
    )
    result["classification"]["total"] = total_cls
    result["detection"]["total_images"] = sum(
        result["detection"][s]["images"] for s in ["train", "val"]
    )

    return result