from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


DRONE_NAMES = ("H80.JPG", "H90.JPG", "H100.JPG")
SAT_NAMES = ("H80.tif", "H90.tif", "H100.tif", "H80_old.tif", "H90_old.tif", "H100_old.tif")


@dataclass(frozen=True)
class DenseUavLayout:
    root: Path

    @property
    def train_drone_dir(self) -> Path:
        return self.root / "train" / "drone"

    @property
    def train_sat_dir(self) -> Path:
        return self.root / "train" / "satellite"

    @property
    def test_query_drone_dir(self) -> Path:
        return self.root / "test" / "query_drone"

    @property
    def test_gallery_sat_dir(self) -> Path:
        return self.root / "test" / "gallery_satellite"


def _iter_ids(dir_path: Path) -> list[str]:
    if not dir_path.exists():
        return []
    ids = [p.name for p in dir_path.iterdir() if p.is_dir()]
    return sorted(ids)


def _collect_files_for_id(id_dir: Path, expected_names: Iterable[str]) -> dict[str, str]:
    """
    Returns mapping name -> relative path (posix) for files that exist.
    """
    out: dict[str, str] = {}
    for name in expected_names:
        p = id_dir / name
        if p.exists():
            out[name] = p.as_posix()
    return out


def _write_lines(path: Path, lines: Iterable[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for line in lines:
            f.write(line.rstrip("\n") + "\n")


def build_indices(
    root: Path,
    out_dir: Path,
    *,
    exclude_old: bool = False,
    strict: bool = True,
    satellite_ext: str = ".tif",
) -> dict:
    """
    Создаёт train/test индексы в стиле:
        query_path label pos1 pos2 ...
    где label = целочисленный id класса (по порядку).

    Positive list:
      - для train query: все доступные satellite варианты того же ID (3 или 6 файлов)
      - для test query: satellite варианты из test/gallery_satellite/<ID> (если есть)
    """
    layout = DenseUavLayout(root=root)

    # IDs
    train_ids = _iter_ids(layout.train_drone_dir)
    train_sat_ids = _iter_ids(layout.train_sat_dir)
    if strict and train_ids != train_sat_ids:
        raise ValueError("train/drone IDs differ from train/satellite IDs")

    gallery_ids = _iter_ids(layout.test_gallery_sat_dir)
    query_ids = _iter_ids(layout.test_query_drone_dir)

    # label mapping: use gallery id universe for stable evaluation labels
    # (train ids are subset of gallery ids in typical setting)
    all_ids = sorted(set(gallery_ids) | set(train_ids) | set(query_ids))
    id_to_label = {id_: i for i, id_ in enumerate(all_ids)}

    # Collect DB lists
    def sat_paths_for(id_: str, sat_root: Path) -> list[str]:
        id_dir = sat_root / id_
        if not id_dir.exists():
            return []
        names = list(SAT_NAMES)
        if exclude_old:
            names = [n for n in names if not n.endswith("_old.tif")]
        # allow processed datasets where tif were converted to png
        names_fs = [n.replace(".tif", satellite_ext) for n in names]
        got = _collect_files_for_id(id_dir, names_fs)
        # stable order: H80,H90,H100,(old...)
        ordered = [got[n] for n in names_fs if n in got]
        # make relative to dataset root
        rel = [str(Path(p).relative_to(root).as_posix()) for p in ordered]
        return rel

    def drone_paths_for(id_: str, drone_root: Path) -> list[str]:
        id_dir = drone_root / id_
        if not id_dir.exists():
            return []
        got = _collect_files_for_id(id_dir, DRONE_NAMES)
        ordered = [got[n] for n in DRONE_NAMES if n in got]
        rel = [str(Path(p).relative_to(root).as_posix()) for p in ordered]
        return rel

    # train_db: all train satellite images (optionally exclude old)
    train_db: list[str] = []
    for id_ in train_ids:
        train_db.extend(sat_paths_for(id_, layout.train_sat_dir))

    # test_db: all test gallery satellite images
    test_db: list[str] = []
    for id_ in gallery_ids:
        test_db.extend(sat_paths_for(id_, layout.test_gallery_sat_dir))

    # Queries:
    train_query_lines: list[str] = []
    train_missing: list[str] = []
    for id_ in train_ids:
        q_paths = drone_paths_for(id_, layout.train_drone_dir)
        pos = sat_paths_for(id_, layout.train_sat_dir)
        if strict:
            if len(q_paths) != 3:
                train_missing.append(f"{id_}: drone files {len(q_paths)}/3")
            if (exclude_old and len(pos) != 3) or ((not exclude_old) and len(pos) != 6):
                train_missing.append(f"{id_}: satellite files {len(pos)}/expected")
        for q in q_paths:
            label = id_to_label[id_]
            if not pos:
                if strict:
                    raise ValueError(f"No positives for train query id={id_}")
                continue
            train_query_lines.append(" ".join([q, str(label), *pos]))

    test_query_lines: list[str] = []
    test_missing: list[str] = []
    for id_ in query_ids:
        q_paths = drone_paths_for(id_, layout.test_query_drone_dir)
        pos = sat_paths_for(id_, layout.test_gallery_sat_dir)
        if strict:
            if len(q_paths) != 3:
                test_missing.append(f"{id_}: query files {len(q_paths)}/3")
            if not pos:
                test_missing.append(f"{id_}: no gallery satellite folder/files")
        for q in q_paths:
            label = id_to_label[id_]
            if not pos:
                if strict:
                    raise ValueError(f"No positives for test query id={id_}")
                continue
            test_query_lines.append(" ".join([q, str(label), *pos]))

    # Write outputs
    index_dir = out_dir / "index"
    _write_lines(index_dir / "train_db.txt", train_db)
    _write_lines(index_dir / "test_db.txt", test_db)
    _write_lines(index_dir / "train_query.txt", train_query_lines)
    _write_lines(index_dir / "test_query.txt", test_query_lines)

    stats = {
        "root": str(root),
        "out_dir": str(out_dir),
        "exclude_old": exclude_old,
        "strict": strict,
        "counts": {
            "n_train_ids": len(train_ids),
            "n_gallery_ids": len(gallery_ids),
            "n_query_ids": len(query_ids),
            "n_all_ids_universe": len(all_ids),
            "n_train_db_images": len(train_db),
            "n_test_db_images": len(test_db),
            "n_train_query_images": len(train_query_lines),
            "n_test_query_images": len(test_query_lines),
        },
        "integrity": {
            "train_issues": train_missing[:200],
            "test_issues": test_missing[:200],
        },
    }
    (out_dir / "stats").mkdir(parents=True, exist_ok=True)
    (out_dir / "stats" / "stats.json").write_text(json.dumps(stats, ensure_ascii=False, indent=2), encoding="utf-8")
    return stats