"""Entry point for the depth/edges/segmentation/chmv2 augmentation pipeline.

All parameters loaded from gin config files — no argparse.
Sequential stage processing: one model at a time, load → process all → unload.

Usage:
    python -m src.main
"""

from __future__ import annotations

import gc
import json
import logging
import time
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from src.conf.config_loader import base_gin_files, load_all_configs
from src.conf.pipeline_conf import PipelineConfig
from src.conf.hardware_conf import HardwareConfig
from src.conf.models_conf import ModelsConfig
from src.conf.input_conf import InputConfig
from src.conf.seg_conf import SegConfig
from src.utils.utils_file_dir import get_proj_dir

from src.augmentor.dataset import (
    AugmentDataset, ImageRecord, attach_output_dirs,
    discover_images, filter_completed, split_by_view,
)
from src.augmentor.inference import (
    compute_edges_from_depth, infer_depth_batch, infer_chmv2_batch,
    infer_segmentation_batch, postprocess_segmentation,
)
from src.augmentor.io_utils import (
    npy_path, vis_path,
    save_depth_async, save_chmv2_async, save_edges_async,
    save_segmentation_async, consolidate_safetensors,
    setup_logging, shutdown_io_pool,
)
from src.augmentor.models import (
    load_depth_model, load_chmv2_model, load_segmentation_model, unload_model,
)
from src.utils.profiler import (
    log_system_info, log_disk_info, log_vram_snapshot, log_ram_snapshot,
)

logger = logging.getLogger(__name__)

_STAGE_EMOJI = {
    "depth": "🌊",
    "edges": "🔪",
    "segmentation": "🗺️",
    "chmv2": "🦕",
    "consolidate": "📦",
}


def _silence_model_loggers() -> None:
    """Suppress verbose inference logs from all models."""
    import os
    import warnings

    os.environ["DA3_LOG_LEVEL"] = "ERROR"
    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

    for name in (
        "depth_anything_3", "depth_anything_3.api",
        "depth_anything_3.utils.logger",
        "transformers", "transformers.modeling_utils",
        "transformers.configuration_utils", "transformers.image_processing_utils",
        "transformers.image_processing_base",
        "sam3", "segearthov3_segmentor",
        "huggingface_hub", "httpx", "filelock",
        "numexpr", "numexpr.utils",
        "torch", "py.warnings",
    ):
        logging.getLogger(name).setLevel(logging.ERROR)

    warnings.filterwarnings("ignore", category=DeprecationWarning)
    warnings.filterwarnings("ignore", category=FutureWarning)
    warnings.filterwarnings("ignore", message=".*not sharded.*")


# ---------------------------------------------------------------------------
# Stage runners
# ---------------------------------------------------------------------------

def _resolve_image_sizes(
    records: list[ImageRecord],
    input_conf: InputConfig,
) -> list[tuple[list[ImageRecord], int, str]]:
    """Split records into groups by target resolution."""
    if input_conf.query_image_size == input_conf.image_size:
        return [(records, input_conf.image_size, "all")]
    db_recs, query_recs = split_by_view(records)
    groups: list[tuple[list[ImageRecord], int, str]] = []
    if db_recs:
        groups.append((db_recs, input_conf.image_size, f"db {input_conf.image_size}"))
    if query_recs:
        groups.append((query_recs, input_conf.query_image_size, f"query {input_conf.query_image_size}"))
    return groups


def run_depth_stage(
    records: list[ImageRecord],
    pipeline_conf: PipelineConfig,
    hw_conf: HardwareConfig,
    models_conf: ModelsConfig,
    input_conf: InputConfig,
    device: torch.device,
) -> None:
    """🌊 Load DA3, process all images, unload."""
    model = load_depth_model(models_conf, hw_conf, device)

    for group_records, sz, label in _resolve_image_sizes(records, input_conf):
        bs = hw_conf.find_batch_size(
            inference_fn=lambda x: infer_depth_batch(model, x, device),
            input_shape=(3, sz, sz),
            device=device,
        )

        ds = AugmentDataset(group_records, image_size=sz)
        logger.info("🌊 [depth/%s] DataLoader: batch_size=%d, %d images, %d batches",
                     label, bs, len(ds), (len(ds) + bs - 1) // bs)
        loader = DataLoader(
            ds, batch_size=bs, shuffle=False,
            num_workers=hw_conf.num_workers, pin_memory=True,
            persistent_workers=hw_conf.num_workers > 0,
            prefetch_factor=4 if hw_conf.num_workers > 0 else None,
        )

        total_images = len(ds)
        pbar = tqdm(loader, desc=f"🌊 depth/{label} (bs={bs})", unit="batch",
                    total=len(loader), colour="green",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
        processed = 0
        for batch in pbar:
            depths = infer_depth_batch(model, batch["image_raw"], device)
            for i in range(depths.shape[0]):
                save_depth_async(depths[i],
                                 Path(batch["output_root"][i]),
                                 batch["rel_parent"][i],
                                 stem=batch["stem"][i],
                                 save_npy=pipeline_conf.save_npy,
                                 save_vis=pipeline_conf.save_vis)
            processed += depths.shape[0]
            pbar.set_postfix(images=f"{processed}/{total_images}")

        shutdown_io_pool()

    unload_model(model)


def _load_depth_tensor(r: ImageRecord) -> torch.Tensor:
    """Load depth from npy or png, return [1, H, W] float32."""
    np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
    vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
    if np_p.exists():
        d = np.load(np_p).astype(np.float32)
    else:
        from PIL import Image as _Img
        d = np.array(_Img.open(vis_p)).astype(np.float32) / 255.0
        if d.ndim == 2:
            d = d[np.newaxis]
    t = torch.from_numpy(d)
    if t.ndim == 2:
        t = t.unsqueeze(0)
    return t


def run_edges_stage(
    records: list[ImageRecord],
    pipeline_conf: PipelineConfig,
    batch_size: int = 32,
) -> None:
    """🔪 Compute Sobel edges from saved depth (CPU, batched)."""
    valid: list[ImageRecord] = []
    for r in records:
        np_p = npy_path(r.output_root, "depth", r.rel_parent, r.stem)
        vis_p = vis_path(r.output_root, "depth", r.rel_parent, r.stem)
        if np_p.exists() or vis_p.exists():
            valid.append(r)
        else:
            logger.warning("⚠️  No depth for %s, skipping edges.", r.rel_path)

    # Group by spatial resolution to avoid stack errors with mixed sizes.
    by_size: dict[tuple[int, int], list[ImageRecord]] = {}
    for r in valid:
        t = _load_depth_tensor(r)
        hw = (t.shape[-2], t.shape[-1])
        by_size.setdefault(hw, []).append(r)

    total_images = len(valid)
    processed = 0
    for (H, W), group in by_size.items():
        label = f"🔪 edges {H}x{W} (sobel)"
        pbar = tqdm(range(0, len(group), batch_size), desc=label, unit="batch",
                    colour="cyan",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
        for start in pbar:
            chunk = group[start : start + batch_size]
            depth_tensors = [_load_depth_tensor(r) for r in chunk]
            depths = torch.stack(depth_tensors)
            if depths.ndim == 3:
                depths = depths.unsqueeze(1)
            edges_batch = compute_edges_from_depth(depths)
            for j, r in enumerate(chunk):
                save_edges_async(edges_batch[j], r.output_root, r.rel_parent,
                                 stem=r.stem,
                                 save_npy=pipeline_conf.save_npy,
                                 save_vis=pipeline_conf.save_vis)
            processed += len(chunk)
            pbar.set_postfix(images=f"{processed}/{total_images}")

    shutdown_io_pool()


def run_chmv2_stage(
    records: list[ImageRecord],
    pipeline_conf: PipelineConfig,
    hw_conf: HardwareConfig,
    models_conf: ModelsConfig,
    input_conf: InputConfig,
    device: torch.device,
) -> None:
    """🦕 Load CHMv2 (DINOv3), process all images, unload."""
    model, processor = load_chmv2_model(models_conf, hw_conf, device)

    for group_records, sz, label in _resolve_image_sizes(records, input_conf):
        bs = hw_conf.find_batch_size(
            inference_fn=lambda x: infer_chmv2_batch(model, processor, x, device),
            input_shape=(3, sz, sz),
            device=device,
        )

        ds = AugmentDataset(group_records, image_size=sz)
        logger.info("🦕 [chmv2/%s] DataLoader: batch_size=%d, %d images, %d batches",
                     label, bs, len(ds), (len(ds) + bs - 1) // bs)
        loader = DataLoader(
            ds, batch_size=bs, shuffle=False,
            num_workers=hw_conf.num_workers, pin_memory=True,
            persistent_workers=hw_conf.num_workers > 0,
            prefetch_factor=4 if hw_conf.num_workers > 0 else None,
        )

        total_images = len(ds)
        pbar = tqdm(loader, desc=f"🦕 chmv2/{label} (bs={bs})", unit="batch",
                    total=len(loader), colour="blue",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
        processed = 0
        for batch in pbar:
            depths = infer_chmv2_batch(model, processor, batch["image_raw"], device)
            for i in range(depths.shape[0]):
                save_chmv2_async(depths[i],
                                 Path(batch["output_root"][i]),
                                 batch["rel_parent"][i],
                                 stem=batch["stem"][i],
                                 save_npy=pipeline_conf.save_npy,
                                 save_vis=pipeline_conf.save_vis)
            processed += depths.shape[0]
            pbar.set_postfix(images=f"{processed}/{total_images}")

        shutdown_io_pool()

    unload_model(model)


def run_segmentation_stage(
    records: list[ImageRecord],
    pipeline_conf: PipelineConfig,
    hw_conf: HardwareConfig,
    models_conf: ModelsConfig,
    input_conf: InputConfig,
    seg_conf: SegConfig,
    device: torch.device,
) -> None:
    """🗺️ Load segmentation model, process all images, unload."""
    model, seg_config = load_segmentation_model(models_conf, hw_conf, seg_conf, device)
    num_classes = seg_config.get("num_classes", 150)
    is_segearth = seg_config.get("type") == "segearth-ov3"

    for group_records, sz, label in _resolve_image_sizes(records, input_conf):
        if is_segearth:
            bs = 16
        else:
            bs = hw_conf.find_batch_size(
                inference_fn=lambda x: infer_segmentation_batch(model, seg_config, x, device),
                input_shape=(3, sz, sz),
                device=device,
            )

        ds = AugmentDataset(group_records, image_size=sz)
        total_images = len(ds)
        logger.info("🗺️  [segmentation/%s] DataLoader: batch_size=%d, %d images, %d batches",
                     label, bs, total_images, (total_images + bs - 1) // bs)
        loader = DataLoader(
            ds, batch_size=bs, shuffle=False,
            num_workers=hw_conf.num_workers, pin_memory=True,
            persistent_workers=hw_conf.num_workers > 0,
            prefetch_factor=4 if hw_conf.num_workers > 0 else None,
        )

        seg_type = "SegEarth-OV3" if is_segearth else "SegFormer"
        pbar = tqdm(loader, desc=f"🗺️  seg/{label} {seg_type} (bs={bs})", unit="batch",
                    colour="yellow",
                    bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} batches [{elapsed}<{remaining}, {rate_fmt}]")
        processed = 0
        for batch in pbar:
            segs = infer_segmentation_batch(
                model, seg_config, batch["image_raw"], device,
            )
            if (pipeline_conf.seg_fix_dark_water
                    or pipeline_conf.seg_fix_uniform_water
                    or pipeline_conf.seg_reclassify_wetland):
                segs = postprocess_segmentation(
                    segs, batch["image_raw"],
                    fix_uniform_water=pipeline_conf.seg_fix_uniform_water,
                    reclassify_wetland=pipeline_conf.seg_reclassify_wetland,
                )
            for j in range(segs.shape[0]):
                save_segmentation_async(
                    segs[j], Path(batch["output_root"][j]),
                    batch["rel_parent"][j],
                    stem=batch["stem"][j],
                    save_npy=pipeline_conf.save_npy, save_vis=pipeline_conf.save_vis,
                    num_classes=num_classes,
                )
            processed += segs.shape[0]
            pbar.set_postfix(images=f"{processed}/{total_images}")

        shutdown_io_pool()

    unload_model(model)


def run_consolidate_stage(
    records: list[ImageRecord],
    cleanup_npy: bool = False,
) -> None:
    """📦 Bundle per-image .npy/.png modalities into .safetensors files."""
    total = len(records)
    written = 0
    pbar = tqdm(records, desc="📦 consolidate → safetensors", unit="img",
                colour="magenta",
                bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]")
    for r in pbar:
        ok = consolidate_safetensors(r.output_root, r.rel_parent, r.stem,
                                     cleanup_npy=cleanup_npy)
        if ok:
            written += 1
        pbar.set_postfix(written=f"{written}/{total}")
    logger.info("📦 Consolidated %d / %d images to .safetensors.", written, total)


# ---------------------------------------------------------------------------
# Pipeline orchestration
# ---------------------------------------------------------------------------

def run_pipeline(
    pipeline_conf: PipelineConfig,
    hw_conf: HardwareConfig,
    models_conf: ModelsConfig,
    input_conf: InputConfig,
    seg_conf: SegConfig,
) -> None:
    """Execute the full augmentation pipeline: one stage at a time."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type != "cuda":
        logger.warning("⚠️  CUDA not available, running on CPU (very slow).")

    _silence_model_loggers()

    # System profiling at startup.
    log_system_info()
    print(Path(pipeline_conf.input_root))
    log_disk_info(Path(pipeline_conf.input_root), Path(pipeline_conf.output_root))

    # Discover images.
    input_root = Path(pipeline_conf.input_root)
    output_root = Path(pipeline_conf.output_root)
    logger.info(
        "🔍 Discovering images in %s (subset=%s, source=%s) ...",
        input_root, pipeline_conf.subset or "all", pipeline_conf.source or "all",
    )
    all_records = discover_images(input_root, subset=pipeline_conf.subset,
                                  source=pipeline_conf.source)
    all_records = attach_output_dirs(all_records, output_root)
    logger.info("📸 Found %d images.", len(all_records))

    if not all_records:
        logger.error("❌ No images found. Check input_root in pipeline.gin.")
        return

    # When save_safetensors is on, always save intermediate .npy
    # (consolidation reads .npy for lossless float16; PNG is lossy uint8).
    if pipeline_conf.save_safetensors and not pipeline_conf.save_npy:
        pipeline_conf.save_npy = True
        if not pipeline_conf.cleanup_npy:
            pipeline_conf.cleanup_npy = True
            logger.info(
                "📦 save_safetensors=True → forcing save_npy=True + cleanup_npy=True "
                "(intermediate .npy will be removed after consolidation)."
            )

    # Process each stage sequentially.
    stage_times: dict[str, float] = {}
    stage_counts: dict[str, int] = {}
    failed_stages: set[str] = set()

    for stage in pipeline_conf.stages:
        emoji = _STAGE_EMOJI.get(stage, "⚙️")

        if stage == "edges" and "depth" in failed_stages:
            logger.error("❌ [edges] Skipped — depth stage failed.")
            failed_stages.add(stage)
            continue

        pending, skipped = filter_completed(all_records, stage)

        logger.info("%s [%s] %d pending, %d skipped.", emoji, stage, len(pending), skipped)
        stage_counts[stage] = len(pending)

        if not pending:
            stage_times[stage] = 0.0
            continue

        logger.info("%s [%s] Starting stage...", emoji, stage)
        log_vram_snapshot(f"before {stage}")
        log_ram_snapshot(f"before {stage}")
        t0 = time.perf_counter()
        try:
            if stage == "depth":
                run_depth_stage(pending, pipeline_conf, hw_conf, models_conf,
                                input_conf, device)
            elif stage == "edges":
                run_edges_stage(pending, pipeline_conf)
            elif stage == "segmentation":
                run_segmentation_stage(pending, pipeline_conf, hw_conf, models_conf,
                                       input_conf, seg_conf, device)
            elif stage == "chmv2":
                run_chmv2_stage(pending, pipeline_conf, hw_conf, models_conf,
                                input_conf, device)
        except Exception:
            logger.exception("💥 Stage '%s' failed.", stage)
            failed_stages.add(stage)
        elapsed = time.perf_counter() - t0
        stage_times[stage] = elapsed
        if stage not in failed_stages:
            logger.info("✅ [%s] Completed in %.1f s (%d images).", stage, elapsed, len(pending))
        log_vram_snapshot(f"after {stage}")
        log_ram_snapshot(f"after {stage}")

    # SafeTensors consolidation: bundle all modalities per image.
    if pipeline_conf.save_safetensors:
        pending_st, skipped_st = filter_completed(all_records, "consolidate")
        logger.info("📦 [consolidate] %d pending, %d skipped.", len(pending_st), skipped_st)
        if pending_st:
            t0 = time.perf_counter()
            run_consolidate_stage(pending_st, cleanup_npy=pipeline_conf.cleanup_npy)
            elapsed_st = time.perf_counter() - t0
            stage_times["consolidate"] = elapsed_st
            stage_counts["consolidate"] = len(pending_st)
            logger.info("✅ [consolidate] Completed in %.1f s (%d images).",
                         elapsed_st, len(pending_st))
        else:
            stage_times["consolidate"] = 0.0
            stage_counts["consolidate"] = 0

    # Manifest.
    manifest = {
        "pipeline_version": "4.0.0-dir-layout",
        "image_size_db": input_conf.image_size,
        "image_size_query": input_conf.query_image_size,
        "profile": hw_conf.profile_name,
        "models": {
            "depth": models_conf.depth_model_id,
            "edges": "Sobel from depth (CPU)",
            "segmentation": models_conf.seg_model_type,
            "chmv2": models_conf.chmv2_model_id,
        },
        "save_safetensors": pipeline_conf.save_safetensors,
        "seg_prompts": seg_conf.prompts,
        "total_images": len(all_records),
        "stages": {
            s: {"processed": stage_counts.get(s, 0),
                "time_sec": round(stage_times.get(s, 0), 1)}
            for s in pipeline_conf.stages
        },
        "timestamp": datetime.now().isoformat(),
    }
    manifest_path = output_root / "manifest.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(
        json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8",
    )

    # Summary.
    total = sum(stage_times.values())
    logger.info("=" * 60)
    logger.info("🏁 DONE: %d images, %.1f s total", len(all_records), total)
    for s, t in stage_times.items():
        cnt = stage_counts.get(s, len(all_records))
        fps = cnt / t if t > 0 else 0
        emoji = _STAGE_EMOJI.get(s, "⚙️")
        logger.info("  %s %-13s %6.1f s  (%d images, %.0f FPS)", emoji, s, t, cnt, fps)
    logger.info("📂 Output: %s", output_root)
    logger.info("=" * 60)


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    """Load all gin configs and run the augmentation pipeline."""
    import argparse

    parser = argparse.ArgumentParser(description="Augmentation pipeline")
    parser.add_argument("--gin", action="append", default=[],
                        help="Gin parameter overrides (repeatable)")
    args = parser.parse_args()

    _silence_model_loggers()

    proj_dir = get_proj_dir()
    path2cfg = f"{proj_dir}in/config_files/"

    if args.gin:
        import gin as _gin
        cfg_dir = Path(path2cfg)
        gin_files = base_gin_files(cfg_dir)
        _gin.clear_config()
        _gin.parse_config_files_and_bindings(
            config_files=[str(f) for f in gin_files],
            bindings=args.gin,
        )
        configs = {
            "pipeline": PipelineConfig(),
            "hardware": HardwareConfig(),
            "models": ModelsConfig(),
            "input": InputConfig(),
            "seg": SegConfig(),
        }
    else:
        configs = load_all_configs(path2cfg)

    pipeline_conf: PipelineConfig = configs["pipeline"]
    setup_logging(pipeline_conf.log_level,
                  log_file=Path(pipeline_conf.output_root) / "pipeline.log")

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

    run_pipeline(
        configs["pipeline"],
        configs["hardware"],
        configs["models"],
        configs["input"],
        configs["seg"],
    )


if __name__ == "__main__":
    main()
