#!/usr/bin/env python3
"""Run annotation pipeline on an arbitrary folder of RGB images or a single image.

Generic entry point: no dataset-specific assumptions. Images are discovered
recursively with a local walker (relative paths preserved in the output
layout); the World-UAV incomplete-scene and service-dir filters of
``discover_images`` are deliberately NOT applied — an arbitrary folder may
legitimately contain directories named ``SoHo`` or ``Index``. A single image
file is also accepted — its modalities land directly under the output root.

Usage:
    python scripts/run_folder.py /path/to/images
    python scripts/run_folder.py /path/to/photo.jpg
    python scripts/run_folder.py /path/to/images --output /path/to/out
    python scripts/run_folder.py /path/to/images --stages depth edges --no-vis
"""

from __future__ import annotations

import argparse
import logging
import sys
from pathlib import Path

# Ensure project root is importable.
_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

import numpy as np
import torch

from src.conf.hardware_conf import HardwareConfig
from src.conf.input_conf import InputConfig
from src.conf.models_conf import ModelsConfig
from src.conf.pipeline_conf import PipelineConfig
from src.conf.seg_conf import SegConfig
from src.augmentor.dataset import EXTENSIONS, ImageRecord
from src.augmentor.io_utils import setup_logging
from src.main import run_pipeline
from scripts.seg_classes import UNIFIED_PROMPTS

logger = logging.getLogger(__name__)

#: Only truly generic service dirs are skipped; the dataset-specific
#: EXCLUDE_DIRS / INCOMPLETE_SCENES of ``discover_images`` do not apply here.
GENERIC_EXCLUDE_DIRS = {"__pycache__", "__MACOSX"}


def discover_folder_images(
    root: Path,
    source: str | None = None,
) -> list[ImageRecord]:
    """Recursively find images under *root* without dataset-specific filters.

    Args:
        root: Folder to scan.
        source: Optional filter — 'query' keeps drone/query images, 'db' keeps
            satellite/DB images (same path-part convention as discover_images).

    Returns:
        Sorted list of ImageRecord with output_root left unset.
    """
    records: list[ImageRecord] = []
    for p in sorted(root.rglob("*")):
        if not p.is_file() or p.suffix.lower() not in EXTENSIONS:
            continue
        if any(d in p.parts for d in GENERIC_EXCLUDE_DIRS):
            continue
        rel = p.relative_to(root)
        if source is not None:
            is_db = "DB" in rel.parts or "satellite" in rel.parts
            is_query = "query" in rel.parts or "drone" in rel.parts
            if source == "query" and is_db:
                continue
            if source == "db" and is_query:
                continue
        records.append(ImageRecord(
            abs_path=p,
            rel_path=str(rel),
            stem=p.stem,
            output_root=Path(),
            rel_parent=str(rel.parent),
        ))
    return records


def build_parser() -> argparse.ArgumentParser:
    """Build the argparse parser (separate function for testability)."""
    parser = argparse.ArgumentParser(
        description="Annotate an arbitrary folder of RGB images or a single image",
    )
    parser.add_argument("input",
                        help="Path to a folder with RGB images or a single "
                             "image file (.png/.jpg/.jpeg/.bmp)")
    parser.add_argument("--output", default=None,
                        help="Output root (default: sibling '<input>-aug'; "
                             "for a single file: '<parent>-aug')")
    parser.add_argument("--stages", nargs="+",
                        default=["depth", "edges", "segmentation", "chmv2"],
                        help="Stages to run")
    parser.add_argument("--image-size", type=int, default=256,
                        help="Output resolution for db/satellite images")
    parser.add_argument("--query-image-size", type=int, default=None,
                        help="Output resolution for query/drone images "
                             "(default: same as --image-size; only relevant "
                             "when the folder contains drone/query subdirs)")
    parser.add_argument("--source", choices=["db", "drone", "all"], default="all",
                        help="Process only db (satellite), drone, or all "
                             "(default; ignored for a single-file input)")
    parser.add_argument("--no-vis", action="store_true",
                        help="Do not save PNG visualizations")
    parser.add_argument("--no-safetensors", action="store_true",
                        help="Do not consolidate modalities into .safetensors")
    parser.add_argument("--wetland-reclassify", action="store_true",
                        help="Reclassify wetland into vegetation/bare soil "
                             "(GTA-specific post-processing, off by default)")
    parser.add_argument("--weights-dir",
                        default=str(_PROJECT_ROOT / "in" / "weights"),
                        help="Directory with model weights")
    parser.add_argument("--num-workers", type=int, default=4,
                        help="DataLoader workers")
    parser.add_argument("--profile", default="rtx4090",
                        help="Hardware profile name")
    return parser


def resolve_output_root(input_path: Path, output: str | None) -> Path:
    """Return the output root: explicit --output or a '-aug' sibling."""
    if output is not None:
        return Path(output)
    base = input_path if input_path.is_dir() else input_path.parent
    return base.parent / f"{base.name}-aug"


def build_single_file_record(image_path: Path) -> ImageRecord:
    """Build one ImageRecord for a standalone image file."""
    return ImageRecord(
        abs_path=image_path,
        rel_path=image_path.name,
        stem=image_path.stem,
        output_root=Path(),
        rel_parent=".",
    )


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.no_vis and args.no_safetensors:
        parser.error(
            "--no-vis together with --no-safetensors would run inference "
            "without writing any output; drop one of the flags"
        )

    import gin
    gin.clear_config()

    input_path = Path(args.input).resolve()
    if not input_path.exists():
        raise SystemExit(f"Input path does not exist: {input_path}")

    source = None if args.source == "all" else args.source
    if source == "drone":
        source = "query"

    if input_path.is_file():
        if input_path.suffix.lower() not in EXTENSIONS:
            raise SystemExit(
                f"Unsupported image extension '{input_path.suffix}' "
                f"(expected one of {sorted(EXTENSIONS)})"
            )
        records = [build_single_file_record(input_path)]
        input_root = input_path.parent
    else:
        input_root = input_path
        records = discover_folder_images(input_root, source=source)
        if not records:
            raise SystemExit(f"No RGB images found under: {input_root}")

    output_root = resolve_output_root(input_path, args.output)

    pipeline_conf = PipelineConfig(
        input_root=str(input_root),
        output_root=str(output_root),
        stages=args.stages,
        save_npy=False,
        save_vis=not args.no_vis,
        save_safetensors=not args.no_safetensors,
        cleanup_npy=True,
        seg_fix_dark_water=True,
        seg_reclassify_wetland=args.wetland_reclassify,
        resume=True,
        source=source,
        log_level="INFO",
    )

    hw_conf = HardwareConfig(
        profile_name=args.profile,
        total_ram_gb=24.0,
        reserve_gb=2.0,
        use_fp16=True,
        batch_size=None,  # auto
        num_workers=args.num_workers,
    )

    input_conf = InputConfig(
        image_size=args.image_size,
        query_image_size=args.query_image_size or args.image_size,
    )

    # Unified 17-class prompt list — shared across all datasets.
    seg_conf = SegConfig(threshold=0.15, prompts=UNIFIED_PROMPTS)

    models_conf = ModelsConfig(weights_dir=args.weights_dir)

    setup_logging(
        pipeline_conf.log_level,
        log_file=output_root / "pipeline.log",
    )

    if input_path.is_file() and source is not None:
        logger.warning("--source is ignored for a single-file input.")
    logger.info("Input: %s (%d image(s)) -> %s",
                input_path, len(records), output_root)

    torch.manual_seed(42)
    np.random.seed(42)

    run_pipeline(pipeline_conf, hw_conf, models_conf, input_conf, seg_conf,
                 records=records)


if __name__ == "__main__":
    main()