from __future__ import annotations

import argparse
import logging
from pathlib import Path

import coloredlogs

from .gps import parse_dense_gps_txt, write_gps_csv
from .indexing import build_indices
from .processing import ProcessingOptions, process_denseuav_to_new_root


def _p(s: str) -> Path:
    return Path(s).expanduser().resolve()


def main() -> None:
    ap = argparse.ArgumentParser(
        prog="uavdense_prepare",
        description="Prepare DenseUAV (UAVDense): optionally build processed dataset + indices + stats",
    )
    ap.add_argument("--root", required=True, help="Path to raw DenseUAV root directory")
    ap.add_argument("--out", required=True, help="Output directory for indices/stats (always created)")
    ap.add_argument(
        "--dst",
        default=None,
        help="If set, creates a new processed dataset root (resized images, satellite tif->png).",
    )
    ap.add_argument("--target-size", type=int, default=256, help="Resize to NxN in --dst mode")
    ap.add_argument(
        "--limit-ids",
        type=int,
        default=None,
        help="Debug option: in --dst mode process only first N IDs from each split",
    )
    ap.add_argument("--exclude-old", action="store_true", help="Exclude *_old.tif from DB/positives")
    ap.add_argument("--no-strict", action="store_true", help="Do not fail on missing files/ids (best-effort)")
    ap.add_argument("--log-level", default="INFO", help="Logging level (DEBUG, INFO, WARNING, ERROR)")
    args = ap.parse_args()

    coloredlogs.install(level=args.log_level.upper(), fmt="%(asctime)s %(levelname)s %(message)s")
    log = logging.getLogger("uavdense_prepare")

    root = _p(args.root)
    out = _p(args.out)
    strict = not args.no_strict
    dst = _p(args.dst) if args.dst else None

    dataset_root_for_index = root
    satellite_ext = ".tif"

    if dst is not None:
        log.info("Building processed dataset at %s", dst)
        opts = ProcessingOptions(target_size=args.target_size, exclude_old=args.exclude_old, limit_ids=args.limit_ids)
        process_denseuav_to_new_root(src_root=root, dst_root=dst, opts=opts, strict=strict)
        dataset_root_for_index = dst
        satellite_ext = ".png"

    # GPS exports (optional files) from the dataset root we index
    gps_dir = out / "gps"
    for name in ("Dense_GPS_ALL.txt", "Dense_GPS_train.txt", "Dense_GPS_test.txt"):
        src = dataset_root_for_index / name
        if src.exists():
            log.info("Parsing GPS file %s", src)
            recs = parse_dense_gps_txt(src)
            out_csv = gps_dir / (name.replace(".txt", ".csv").lower())
            write_gps_csv(recs, out_csv)

    log.info("Building indices from %s into %s", dataset_root_for_index, out)
    build_indices(
        dataset_root_for_index,
        out,
        exclude_old=args.exclude_old,
        strict=strict,
        satellite_ext=satellite_ext,
    )
    log.info("Done")


if __name__ == "__main__":
    main()