from __future__ import annotations

import json
import re
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable

from fastapi import UploadFile

from ...config import settings

DATASET_KINDS = ("images", "labels", "masks")
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}
LABEL_EXTS = {".txt", ".json", ".yaml", ".yml"}


def uploads_root() -> Path:
    root = settings.project_root / "var" / "uploads" / "datasets"
    root.mkdir(parents=True, exist_ok=True)
    return root


def slugify(value: str) -> str:
    text = re.sub(r"[^A-Za-z0-9_.\-\u4e00-\u9fff]+", "_", value.strip())
    return text.strip("._") or "dataset"


def safe_filename(value: str | None) -> str:
    original = Path(value or "upload.bin").name
    suffix = Path(original).suffix.lower()
    stem = slugify(Path(original).stem or "upload")
    if suffix and re.fullmatch(r"\.[A-Za-z0-9]{1,12}", suffix):
        return f"{stem}{suffix}"
    return stem


def dataset_dir(name: str) -> Path:
    return uploads_root() / slugify(name)


def metadata_path(name: str) -> Path:
    return dataset_dir(name) / "dataset.json"


def create_dataset(name: str, description: str = "") -> dict:
    safe_name = slugify(name)
    root = dataset_dir(safe_name)
    for kind in DATASET_KINDS:
        (root / kind).mkdir(parents=True, exist_ok=True)
    meta = {
        "name": safe_name,
        "description": description,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "root": str(root.relative_to(settings.project_root)),
        "layout": {
            "images": str((root / "images").relative_to(settings.project_root)),
            "labels": str((root / "labels").relative_to(settings.project_root)),
            "masks": str((root / "masks").relative_to(settings.project_root)),
        },
    }
    metadata_path(safe_name).write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    return describe_dataset(safe_name)


def _load_meta(name: str) -> dict:
    path = metadata_path(name)
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    return {"name": slugify(name), "description": "", "root": str(dataset_dir(name).relative_to(settings.project_root))}


def _iter_files(root: Path) -> Iterable[Path]:
    if not root.exists():
        return []
    return sorted(path for path in root.rglob("*") if path.is_file())


def _stem_map(paths: Iterable[Path], exts: set[str] | None = None) -> dict[str, Path]:
    result: dict[str, Path] = {}
    for path in paths:
        if exts and path.suffix.lower() not in exts:
            continue
        result.setdefault(path.stem, path)
    return result


def _image_shape(path: Path) -> dict | None:
    try:
        import cv2

        image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
        if image is None:
            return None
        height, width = image.shape[:2]
        channels = 1 if image.ndim == 2 else image.shape[2]
        return {"width": int(width), "height": int(height), "channels": int(channels)}
    except Exception:
        return None


def _validate_yolo_txt(path: Path) -> dict:
    errors = []
    classes: set[int] = set()
    lines = path.read_text(encoding="utf-8", errors="replace").splitlines()
    annotation_count = 0
    for line_number, line in enumerate(lines, 1):
        raw = line.strip()
        if not raw:
            continue
        parts = raw.split()
        annotation_count += 1
        if len(parts) < 5:
            errors.append(f"{path.name}:{line_number} has fewer than 5 tokens")
            continue
        try:
            class_id = int(float(parts[0]))
            coords = [float(item) for item in parts[1:]]
        except ValueError:
            errors.append(f"{path.name}:{line_number} contains non-numeric values")
            continue
        if class_id < 0:
            errors.append(f"{path.name}:{line_number} has negative class id")
        classes.add(class_id)
        if len(coords) not in {4} and len(coords) < 6:
            errors.append(f"{path.name}:{line_number} has too few segmentation coordinates")
        if len(coords) != 4 and len(coords) % 2 != 0:
            errors.append(f"{path.name}:{line_number} has an odd number of polygon coordinates")
        out_of_range = [value for value in coords if value < 0 or value > 1]
        if out_of_range:
            errors.append(f"{path.name}:{line_number} has coordinates outside 0..1")
    return {"annotations": annotation_count, "classes": sorted(classes), "errors": errors}


def describe_dataset(name: str) -> dict:
    safe_name = slugify(name)
    root = dataset_dir(safe_name)
    meta = _load_meta(safe_name)
    absolute_layout = {kind: str((root / kind).resolve()) for kind in DATASET_KINDS}
    counts = {}
    samples = {}
    for kind in sorted(DATASET_KINDS):
        files = list(_iter_files(root / kind))
        counts[kind] = len(files)
        samples[kind] = [
            {
                "name": path.name,
                "path": str(path.resolve()),
                "relative_path": str(path.resolve().relative_to(settings.project_root)),
                "size": path.stat().st_size,
                "previewable": path.suffix.lower() in IMAGE_EXTS,
            }
            for path in files[:80]
        ]
    return {**meta, "absolute_layout": absolute_layout, "counts": counts, "samples": samples}


def validate_dataset(name: str) -> dict:
    safe_name = slugify(name)
    root = dataset_dir(safe_name)
    image_files = list(_iter_files(root / "images"))
    label_files = list(_iter_files(root / "labels"))
    mask_files = list(_iter_files(root / "masks"))
    images = _stem_map(image_files, IMAGE_EXTS)
    labels = _stem_map(label_files, LABEL_EXTS)
    masks = _stem_map(mask_files, IMAGE_EXTS)

    image_stems = set(images)
    label_stems = set(labels)
    mask_stems = set(masks)
    paired_label_stems = sorted(image_stems & label_stems)
    paired_mask_stems = sorted(image_stems & mask_stems)

    yolo_errors = []
    class_ids: set[int] = set()
    annotation_count = 0
    for label in labels.values():
        if label.suffix.lower() != ".txt":
            continue
        detail = _validate_yolo_txt(label)
        yolo_errors.extend(detail["errors"])
        class_ids.update(detail["classes"])
        annotation_count += detail["annotations"]

    shape_mismatches = []
    sample_shapes = []
    for stem in paired_mask_stems[:80]:
        image_shape = _image_shape(images[stem])
        mask_shape = _image_shape(masks[stem])
        if image_shape and len(sample_shapes) < 8:
            sample_shapes.append({"name": images[stem].name, **image_shape})
        if image_shape and mask_shape and (image_shape["width"], image_shape["height"]) != (mask_shape["width"], mask_shape["height"]):
            shape_mismatches.append({"stem": stem, "image": image_shape, "mask": mask_shape})
    if not sample_shapes:
        for image in list(images.values())[:8]:
            image_shape = _image_shape(image)
            if image_shape:
                sample_shapes.append({"name": image.name, **image_shape})

    checks = [
        {"name": "has_images", "passed": len(images) > 0, "count": len(images)},
        {"name": "has_labels_or_masks", "passed": len(labels) > 0 or len(masks) > 0, "labels": len(labels), "masks": len(masks)},
        {"name": "image_label_pairs", "passed": len(label_stems) == 0 or len(paired_label_stems) > 0, "count": len(paired_label_stems)},
        {"name": "image_mask_pairs", "passed": len(mask_stems) == 0 or len(paired_mask_stems) > 0, "count": len(paired_mask_stems)},
        {"name": "yolo_txt_valid", "passed": not yolo_errors, "errors": yolo_errors[:20]},
        {"name": "mask_shapes_match", "passed": not shape_mismatches, "errors": shape_mismatches[:20]},
    ]
    yolo_ready = len(images) > 0 and len(paired_label_stems) > 0 and not yolo_errors
    mask_ready = len(images) > 0 and len(paired_mask_stems) > 0 and not shape_mismatches
    return {
        "dataset": safe_name,
        "root": str(root.resolve()),
        "counts": {"images": len(images), "labels": len(labels), "masks": len(masks), "annotations": annotation_count},
        "pairs": {
            "image_label": len(paired_label_stems),
            "image_mask": len(paired_mask_stems),
            "images_without_labels": sorted(image_stems - label_stems)[:50],
            "labels_without_images": sorted(label_stems - image_stems)[:50],
            "images_without_masks": sorted(image_stems - mask_stems)[:50],
            "masks_without_images": sorted(mask_stems - image_stems)[:50],
        },
        "classes": sorted(class_ids),
        "sample_shapes": sample_shapes,
        "checks": checks,
        "ready": {"yolo": yolo_ready, "mask": mask_ready, "any": yolo_ready or mask_ready},
    }


def generate_yolo_dataset_yaml(name: str, class_names: list[str] | None = None) -> dict:
    validation = validate_dataset(name)
    if not validation["ready"]["yolo"]:
        raise ValueError("dataset is not YOLO-ready; upload matching images and valid txt labels first")
    safe_name = slugify(name)
    root = dataset_dir(safe_name)
    classes = validation["classes"] or [0]
    class_count = max(classes) + 1
    names = list(class_names or [])
    if len(names) < class_count:
        names.extend(f"class_{index}" for index in range(len(names), class_count))
    yaml_path = root / "dataset.yaml"
    names_block = "\n".join(f"  {index}: {label}" for index, label in enumerate(names))
    yaml_text = "\n".join(
        [
            f"path: {root.resolve()}",
            "train: images",
            "val: images",
            f"nc: {len(names)}",
            "names:",
            names_block,
            "",
        ]
    )
    yaml_path.write_text(yaml_text, encoding="utf-8")
    return {
        "dataset": safe_name,
        "path": str(yaml_path.resolve()),
        "relative_path": str(yaml_path.resolve().relative_to(settings.project_root)),
        "classes": classes,
        "names": names,
        "content": yaml_text,
        "validation": validation,
    }


def list_uploaded_datasets() -> list[dict]:
    root = uploads_root()
    datasets = []
    for item in sorted(root.iterdir()):
        if item.is_dir():
            datasets.append(describe_dataset(item.name))
    return datasets


async def save_upload(dataset: str, kind: str, files: list[UploadFile]) -> dict:
    if kind not in DATASET_KINDS:
        raise ValueError(f"unsupported dataset file kind: {kind}")
    safe_name = slugify(dataset)
    if not metadata_path(safe_name).exists():
        create_dataset(safe_name)
    target = dataset_dir(safe_name) / kind
    target.mkdir(parents=True, exist_ok=True)
    saved = []
    for upload in files:
        filename = safe_filename(upload.filename)
        dst = target / filename
        if dst.exists():
            stem = dst.stem
            suffix = dst.suffix
            counter = 1
            while dst.exists():
                dst = target / f"{stem}_{counter}{suffix}"
                counter += 1
        with dst.open("wb") as handle:
            shutil.copyfileobj(upload.file, handle)
        saved.append({"name": dst.name, "relative_path": str(dst.relative_to(settings.project_root)), "size": dst.stat().st_size})
    return {"dataset": describe_dataset(safe_name), "saved": saved}