"""Rename detection result files so frame stems match ground-truth stems.

This script is intended as a preprocessing step before running the evaluation
pipeline. The evaluator pairs detection and ground-truth files strictly by
their filename stem, so detections that omit a timestamp suffix will be skipped
unless they are renamed to the GT stem.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
from collections import defaultdict
from concurrent.futures import ProcessPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import yaml


DEFAULT_TIMESTAMP_PATTERNS = (
    re.compile(r"(?:[_-]?ts)?\d{10,20}$"),
    re.compile(r"[_-]?\d{8}[_-]\d{6}(?:\d+)?$"),
    re.compile(r"[_-]?\d{8}T\d{6}(?:\d+)?$"),
)
DEFAULT_TIMESTAMP_TOKEN_PATTERN = re.compile(r"^(?:ts)?\d{8,20}$")


@dataclass(frozen=True)
class CaseInfo:
    """One evaluation case directory."""

    det_case_dir: Path
    case_name: str
    level1_name: str | None = None


@dataclass(frozen=True)
class MatchPair:
    """Mapping from one detection file to one GT file."""

    det_file: Path
    gt_file: Path
    strategy: str


@dataclass(frozen=True)
class CaseProcessResult:
    """Worker result for one case."""

    summary: dict[str, object]
    manifest_rows: list[dict[str, object]]


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description=(
            "Rename or copy detection result files so their stems match GT stems "
            "before running eval_tools/core/eval.py."
        )
    )
    parser.add_argument("--config", type=Path, help="Evaluation config YAML used to infer det/gt paths.")
    parser.add_argument("--det-root", type=Path, help="Detection root directory. Overrides config dataset.det_path.")
    parser.add_argument("--gt-root", type=Path, help="Ground-truth root directory. Overrides config dataset.gt_path.")
    parser.add_argument(
        "--path-depth",
        type=int,
        choices=(1, 2),
        help="Directory depth: 1 = root/case, 2 = root/level1/case. Overrides config dataset.path_depth.",
    )
    parser.add_argument(
        "--det-format",
        choices=("auto", "json", "txt"),
        help="Detection format. Overrides config dataset.det_format.",
    )
    parser.add_argument(
        "--gt-format",
        choices=("auto", "json", "txt"),
        help="GT format. Overrides config dataset.gt_format.",
    )
    parser.add_argument(
        "--det-subdir",
        help="Relative detection subdirectory inside each case, e.g. predictions/roi0.",
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        help=(
            "Write renamed files to a mirrored directory tree instead of changing the source det root in place. "
            "Recommended for the first run."
        ),
    )
    parser.add_argument(
        "--in-place",
        action="store_true",
        help="Rename files directly inside the source detection directory.",
    )
    parser.add_argument(
        "--enable-index-fallback",
        action="store_true",
        help=(
            "If normalized-name matching still leaves unmatched files, allow one-to-one pairing by sorted index. "
            "Use only when detection and GT frame order is guaranteed to be aligned."
        ),
    )
    parser.add_argument(
        "--strip-pattern",
        action="append",
        default=[],
        help=(
            "Custom regex used to strip a suffix from stems before matching. Can be specified multiple times. "
            "If omitted, built-in timestamp suffix patterns are used."
        ),
    )
    parser.add_argument(
        "--case",
        dest="cases",
        action="append",
        default=[],
        help="Limit processing to specific case names. Can be specified multiple times.",
    )
    parser.add_argument(
        "--manifest-path",
        type=Path,
        help="Where to save the rename manifest JSON. Defaults under output_root or det_root.",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Return a non-zero exit code if any case has unmatched or ambiguous files.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Preview the mapping without copying or renaming files.",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=1,
        help="Number of case-level worker processes. Use 0 to auto-select based on CPU count.",
    )
    return parser.parse_args()


def load_config(config_path: Path | None) -> dict:
    """Load evaluation YAML config if provided."""
    if config_path is None:
        return {}
    with config_path.open("r", encoding="utf-8") as handle:
        return yaml.safe_load(handle) or {}


def resolve_runtime_options(args: argparse.Namespace, config: dict) -> dict:
    """Resolve CLI options with config fallback."""
    dataset_cfg = config.get("dataset", {}) if isinstance(config, dict) else {}

    det_root = Path(args.det_root or dataset_cfg.get("det_path", "")).expanduser() if (args.det_root or dataset_cfg.get("det_path")) else None
    gt_root = Path(args.gt_root or dataset_cfg.get("gt_path", "")).expanduser() if (args.gt_root or dataset_cfg.get("gt_path")) else None
    path_depth = args.path_depth or dataset_cfg.get("path_depth", 1)
    det_format = args.det_format or dataset_cfg.get("det_format", "auto")
    gt_format = args.gt_format or dataset_cfg.get("gt_format", "auto")
    det_subdir = args.det_subdir if args.det_subdir is not None else dataset_cfg.get("det_subdir")

    if det_root is None or gt_root is None:
        raise ValueError("Both detection root and GT root must be provided by --config or explicit arguments.")
    if args.in_place and args.output_root:
        raise ValueError("Use either --in-place or --output-root, not both.")
    if not args.in_place and args.output_root is None and not args.dry_run:
        raise ValueError("Choose --output-root for a copied result tree or --in-place for source renaming.")

    return {
        "det_root": det_root.resolve(),
        "gt_root": gt_root.resolve(),
        "path_depth": int(path_depth),
        "det_format": det_format,
        "gt_format": gt_format,
        "det_subdir": det_subdir,
    }


def iter_case_dirs(det_root: Path, path_depth: int, selected_cases: set[str]) -> list[CaseInfo]:
    """Enumerate detection case directories according to evaluator path depth."""
    def is_case_dir(path: Path) -> bool:
        return path.is_dir() and not path.name.startswith((".", "_"))

    cases: list[CaseInfo] = []
    if path_depth == 1:
        for case_dir in sorted(p for p in det_root.iterdir() if is_case_dir(p)):
            if selected_cases and case_dir.name not in selected_cases:
                continue
            cases.append(CaseInfo(det_case_dir=case_dir, case_name=case_dir.name))
        return cases

    for level1_dir in sorted(p for p in det_root.iterdir() if is_case_dir(p)):
        for case_dir in sorted(p for p in level1_dir.iterdir() if is_case_dir(p)):
            if selected_cases and case_dir.name not in selected_cases:
                continue
            cases.append(CaseInfo(det_case_dir=case_dir, case_name=case_dir.name, level1_name=level1_dir.name))
    return cases


def resolve_gt_case_dir(gt_root: Path, case_info: CaseInfo, path_depth: int) -> Path:
    """Resolve ground-truth case directory from case info."""
    if path_depth == 1:
        return gt_root / case_info.case_name
    return gt_root / str(case_info.level1_name) / case_info.case_name


def resolve_detection_dir(case_dir: Path, det_format: str, det_subdir: str | None) -> tuple[Path | None, str | None]:
    """Resolve the directory and suffix for detection result files."""
    json_candidates: list[Path] = []
    if det_subdir:
        subdir_path = Path(det_subdir)
        json_candidates.append(subdir_path if subdir_path.is_absolute() else case_dir / subdir_path)
    json_candidates.extend([case_dir / "json_results", case_dir / "predictions"])
    txt_candidate = case_dir / "txt_results"

    def pick_json_dir() -> Path | None:
        for candidate in json_candidates:
            if candidate.exists() and any(candidate.glob("*.json")):
                return candidate
        return None

    if det_format == "json":
        result_dir = pick_json_dir()
        return result_dir, ".json"
    if det_format == "txt":
        return (txt_candidate, ".txt") if txt_candidate.exists() else (None, None)

    result_dir = pick_json_dir()
    if result_dir is not None:
        return result_dir, ".json"
    if txt_candidate.exists() and any(txt_candidate.glob("*.txt")):
        return txt_candidate, ".txt"
    return None, None


def resolve_gt_dir(case_dir: Path, gt_format: str) -> tuple[Path | None, str | None]:
    """Resolve the directory and suffix for GT files."""
    json_candidate = case_dir / "labels_json"
    txt_candidate = case_dir / "labels"

    if gt_format == "json":
        return (json_candidate, ".json") if json_candidate.exists() else (None, None)
    if gt_format == "txt":
        return (txt_candidate, ".txt") if txt_candidate.exists() else (None, None)
    if json_candidate.exists():
        return json_candidate, ".json"
    if txt_candidate.exists():
        return txt_candidate, ".txt"
    return None, None


def build_normalizer(custom_patterns: list[str]) -> Callable[[str], str]:
    """Build a stem normalizer that removes timestamp-like suffixes."""
    patterns = [re.compile(pattern) for pattern in custom_patterns] if custom_patterns else list(DEFAULT_TIMESTAMP_PATTERNS)

    def normalize(stem: str) -> str:
        normalized = stem
        for pattern in patterns:
            normalized = pattern.sub("", normalized)
        tokens = [token for token in normalized.rstrip("_-").split("_") if token]
        filtered_tokens = [token for token in tokens if not DEFAULT_TIMESTAMP_TOKEN_PATTERN.fullmatch(token)]
        return "_".join(filtered_tokens)

    return normalize


def build_pairs(
    det_files: list[Path],
    gt_files: list[Path],
    normalize_stem: Callable[[str], str],
    enable_index_fallback: bool,
) -> tuple[list[MatchPair], list[str], list[dict[str, object]]]:
    """Match detection files to GT files using exact, normalized, then optional index fallback."""
    gt_by_stem = {path.stem: path for path in gt_files}
    used_gt_stems: set[str] = set()
    pairs: list[MatchPair] = []
    issues: list[str] = []
    ambiguous_groups: list[dict[str, object]] = []

    unmatched_det: list[Path] = []
    for det_file in det_files:
        gt_file = gt_by_stem.get(det_file.stem)
        if gt_file is None:
            unmatched_det.append(det_file)
            continue
        used_gt_stems.add(gt_file.stem)
        pairs.append(MatchPair(det_file=det_file, gt_file=gt_file, strategy="exact"))

    unmatched_gt = [gt_file for gt_file in gt_files if gt_file.stem not in used_gt_stems]
    if not unmatched_det or not unmatched_gt:
        return pairs, issues, ambiguous_groups

    det_groups: dict[str, list[Path]] = defaultdict(list)
    gt_groups: dict[str, list[Path]] = defaultdict(list)
    for det_file in unmatched_det:
        det_groups[normalize_stem(det_file.stem)].append(det_file)
    for gt_file in unmatched_gt:
        gt_groups[normalize_stem(gt_file.stem)].append(gt_file)

    remaining_det: list[Path] = []
    remaining_gt: list[Path] = []
    group_keys = sorted(set(det_groups) | set(gt_groups))
    for key in group_keys:
        det_group = sorted(det_groups.get(key, []))
        gt_group = sorted(gt_groups.get(key, []))
        if not det_group:
            remaining_gt.extend(gt_group)
            continue
        if not gt_group:
            remaining_det.extend(det_group)
            continue
        if len(det_group) == 1 and len(gt_group) == 1:
            pairs.append(MatchPair(det_file=det_group[0], gt_file=gt_group[0], strategy="normalized"))
            continue

        ambiguous_groups.append(
            {
                "normalized_key": key,
                "det_files": [path.name for path in det_group],
                "gt_files": [path.name for path in gt_group],
            }
        )
        remaining_det.extend(det_group)
        remaining_gt.extend(gt_group)

    if remaining_det or remaining_gt:
        if enable_index_fallback and len(remaining_det) == len(remaining_gt):
            for det_file, gt_file in zip(sorted(remaining_det), sorted(remaining_gt)):
                pairs.append(MatchPair(det_file=det_file, gt_file=gt_file, strategy="index"))
        else:
            if remaining_det:
                issues.append(
                    "Unmatched detection files: " + ", ".join(path.name for path in sorted(remaining_det))
                )
            if remaining_gt:
                issues.append(
                    "Unused GT files: " + ", ".join(path.name for path in sorted(remaining_gt))
                )

    return pairs, issues, ambiguous_groups


def ensure_parent(path: Path) -> None:
    """Create parent directory if missing."""
    path.parent.mkdir(parents=True, exist_ok=True)


def apply_mapping(
    pairs: list[MatchPair],
    det_root: Path,
    output_root: Path | None,
    in_place: bool,
    dry_run: bool,
) -> list[dict[str, object]]:
    """Apply renaming or copying for matched pairs."""
    manifest_rows: list[dict[str, object]] = []

    for pair in pairs:
        relative_det_path = pair.det_file.relative_to(det_root)
        target_name = f"{pair.gt_file.stem}{pair.det_file.suffix}"
        if output_root is not None:
            target_dir = output_root / relative_det_path.parent
            target_path = target_dir / target_name
            action = "copy"
        else:
            target_path = pair.det_file.with_name(target_name)
            action = "rename"

        row = {
            "strategy": pair.strategy,
            "action": action,
            "source": str(pair.det_file),
            "target": str(target_path),
            "gt_file": str(pair.gt_file),
            "changed": pair.det_file.name != target_path.name,
        }
        manifest_rows.append(row)

        if dry_run:
            continue

        ensure_parent(target_path)
        if output_root is not None:
            if pair.det_file.resolve() == target_path.resolve():
                continue
            shutil.copy2(pair.det_file, target_path)
        elif in_place and pair.det_file != target_path:
            if target_path.exists():
                raise FileExistsError(f"Refusing to overwrite existing file: {target_path}")
            pair.det_file.rename(target_path)

    return manifest_rows


def default_manifest_path(args: argparse.Namespace, runtime: dict) -> Path:
    """Choose a default manifest path."""
    if args.manifest_path is not None:
        return args.manifest_path
    root = args.output_root.resolve() if args.output_root else runtime["det_root"]
    return root / "rename_detection_frames_manifest.json"


def resolve_num_workers(requested_workers: int, case_count: int) -> int:
    """Resolve worker count from CLI input and case count."""
    if case_count <= 1:
        return 1
    if requested_workers <= 0:
        requested_workers = os.cpu_count() or 1
    return max(1, min(requested_workers, case_count))


def process_case(
    case_info: CaseInfo,
    runtime: dict[str, object],
    strip_patterns: list[str],
    enable_index_fallback: bool,
    output_root: str | None,
    in_place: bool,
    dry_run: bool,
) -> CaseProcessResult:
    """Process one case, including match generation and optional file operations."""
    normalize_stem = build_normalizer(strip_patterns)
    det_root = Path(str(runtime["det_root"]))
    gt_root = Path(str(runtime["gt_root"]))
    path_depth = int(runtime["path_depth"])
    det_format = str(runtime["det_format"])
    gt_format = str(runtime["gt_format"])
    det_subdir = runtime.get("det_subdir")
    det_subdir_str = None if det_subdir is None else str(det_subdir)

    gt_case_dir = resolve_gt_case_dir(gt_root, case_info, path_depth)
    det_dir, det_suffix = resolve_detection_dir(case_info.det_case_dir, det_format, det_subdir_str)
    gt_dir, gt_suffix = resolve_gt_dir(gt_case_dir, gt_format)

    summary = {
        "case": case_info.case_name,
        "level1_name": case_info.level1_name,
        "det_case_dir": str(case_info.det_case_dir),
        "gt_case_dir": str(gt_case_dir),
        "status": "ok",
        "issues": [],
        "ambiguous_groups": [],
        "matched": 0,
        "changed": 0,
    }

    if det_dir is None or det_suffix is None:
        summary["status"] = "missing_detection_dir"
        summary["issues"].append("Detection results directory not found or empty.")
        return CaseProcessResult(summary=summary, manifest_rows=[])
    if gt_dir is None or gt_suffix is None:
        summary["status"] = "missing_gt_dir"
        summary["issues"].append("GT labels directory not found.")
        return CaseProcessResult(summary=summary, manifest_rows=[])

    det_files = sorted(path for path in det_dir.iterdir() if path.suffix == det_suffix)
    gt_files = sorted(path for path in gt_dir.iterdir() if path.suffix == gt_suffix)
    if not det_files:
        summary["status"] = "empty_detection_dir"
        summary["issues"].append("No detection files found in resolved directory.")
        return CaseProcessResult(summary=summary, manifest_rows=[])
    if not gt_files:
        summary["status"] = "empty_gt_dir"
        summary["issues"].append("No GT files found in resolved directory.")
        return CaseProcessResult(summary=summary, manifest_rows=[])

    pairs, issues, ambiguous_groups = build_pairs(
        det_files=det_files,
        gt_files=gt_files,
        normalize_stem=normalize_stem,
        enable_index_fallback=enable_index_fallback,
    )
    manifest_rows = apply_mapping(
        pairs=pairs,
        det_root=det_root,
        output_root=Path(output_root) if output_root else None,
        in_place=in_place,
        dry_run=dry_run,
    )

    summary["issues"] = issues
    summary["ambiguous_groups"] = ambiguous_groups
    summary["matched"] = len(pairs)
    summary["changed"] = sum(1 for row in manifest_rows if row["changed"])
    if issues or ambiguous_groups:
        summary["status"] = "partial"
    return CaseProcessResult(summary=summary, manifest_rows=manifest_rows)


def main() -> int:
    """Run the rename workflow."""
    args = parse_args()
    config = load_config(args.config)
    runtime = resolve_runtime_options(args, config)
    selected_cases = set(args.cases)

    case_infos = iter_case_dirs(runtime["det_root"], runtime["path_depth"], selected_cases)
    if not case_infos:
        raise FileNotFoundError(f"No detection cases found under {runtime['det_root']}")
    num_workers = resolve_num_workers(args.num_workers, len(case_infos))

    all_manifest_rows: list[dict[str, object]] = []
    case_summaries: list[dict[str, object]] = []
    failed_cases = 0

    output_root = str(args.output_root.resolve()) if args.output_root else None
    if num_workers == 1:
        results = [
            process_case(
                case_info=case_info,
                runtime=runtime,
                strip_patterns=args.strip_pattern,
                enable_index_fallback=args.enable_index_fallback,
                output_root=output_root,
                in_place=args.in_place,
                dry_run=args.dry_run,
            )
            for case_info in case_infos
        ]
    else:
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            results = list(
                executor.map(
                    process_case,
                    case_infos,
                    [runtime] * len(case_infos),
                    [args.strip_pattern] * len(case_infos),
                    [args.enable_index_fallback] * len(case_infos),
                    [output_root] * len(case_infos),
                    [args.in_place] * len(case_infos),
                    [args.dry_run] * len(case_infos),
                )
            )

    for result in results:
        case_summaries.append(result.summary)
        all_manifest_rows.extend(result.manifest_rows)
        if result.summary["status"] != "ok":
            failed_cases += 1

    manifest = {
        "dry_run": args.dry_run,
        "in_place": args.in_place,
        "output_root": str(args.output_root.resolve()) if args.output_root else None,
        "det_root": str(runtime["det_root"]),
        "gt_root": str(runtime["gt_root"]),
        "path_depth": runtime["path_depth"],
        "det_format": runtime["det_format"],
        "gt_format": runtime["gt_format"],
        "det_subdir": runtime["det_subdir"],
        "cases": case_summaries,
        "mappings": all_manifest_rows,
    }

    manifest_path = default_manifest_path(args, runtime)
    ensure_parent(manifest_path)
    with manifest_path.open("w", encoding="utf-8") as handle:
        json.dump(manifest, handle, ensure_ascii=False, indent=2)

    total_cases = len(case_summaries)
    total_matched = sum(item["matched"] for item in case_summaries)
    total_changed = sum(item["changed"] for item in case_summaries)
    print(f"Processed cases: {total_cases}")
    print(f"Workers used: {num_workers}")
    print(f"Matched files: {total_matched}")
    print(f"Renamed/copied files with changed stems: {total_changed}")
    print(f"Manifest: {manifest_path}")

    if args.strict and failed_cases:
        print(f"Strict mode: {failed_cases} case(s) still have unresolved issues.")
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())