"""Dataset adapter skeletons for external/open music corpora."""

from __future__ import annotations

from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Dict, List
import argparse
import json
import subprocess
import torch
import yaml


AUDIO_EXTS = (".wav", ".mp3", ".flac", ".ogg")
MIN_SMOKE_AUDIO_FILES = 2
MIN_SMOKE_ELIGIBLE_QUERY_FILES = 2


def resolve_device(device: str) -> str:
    if device == "auto":
        return "cuda" if torch.cuda.is_available() else "cpu"
    return device


def load_default_training_config(config_path: str = "configs/default.yaml") -> Dict:
    with open(config_path) as f:
        return yaml.safe_load(f)


def build_smoke_config_summary(
    dataset: str,
    manifests_dir: Path,
    manifest_query_duration: float,
    train_epochs: int,
    batch_size: int,
    requested_device: str,
    resolved_device: str,
    base_cfg: Dict,
) -> Dict:
    return {
        "model": {
            "embed_dim": base_cfg["model"]["embed_dim"],
            "channels": base_cfg["model"]["channels"],
            "n_mels": base_cfg["model"]["n_mels"],
            "use_band_split": base_cfg["model"].get("use_band_split", True),
        },
        "data": {
            "source_dataset": dataset,
            "manifests_dir": str(manifests_dir),
            "manifest_query_duration": manifest_query_duration,
            "train_segment_duration": base_cfg["data"]["segment_dur"],
            "sample_rate": base_cfg["data"]["sample_rate"],
            "n_fft": base_cfg["data"]["n_fft"],
            "hop_length": base_cfg["data"]["hop_length"],
            "query_duration_legacy": manifest_query_duration,
        },
        "run": {
            "train_epochs": train_epochs,
            "batch_size": batch_size,
            "requested_device": requested_device,
            "resolved_device": resolved_device,
        },
    }


@dataclass
class DatasetRecord:
    name: str
    source_url: str
    license: str
    commercial_use: str
    notes: str


class BaseAdapter:
    name = "base"

    def describe(self) -> Dict:
        raise NotImplementedError

    def init_layout(self, root: Path) -> Dict:
        root.mkdir(parents=True, exist_ok=True)
        for sub in ["raw", "processed", "manifests", "licenses"]:
            (root / sub).mkdir(exist_ok=True)
        manifest = {
            "dataset": self.name,
            "root": str(root),
            "status": "initialized",
            "next_steps": [
                "download raw audio according to upstream license terms",
                "convert to catalog/query manifests",
                "record license evidence before training",
            ],
        }
        with open(root / "manifests" / "bootstrap.json", "w") as f:
            json.dump(manifest, f, indent=2, ensure_ascii=False)
        return manifest

    def prepare_local_audio(
        self,
        input_dir: Path,
        output_root: Path,
        eval_ratio: float = 0.2,
        query_duration: float = 8.0,
        query_stride: float | None = None,
        query_strategy: str = "random",
        silence_top_db: int = 30,
        seed: int = 42,
    ) -> Dict:
        output_root.mkdir(parents=True, exist_ok=True)
        cmd = [
            "/usr/local/miniconda3/bin/python",
            "src/data/manifest_tools.py",
            "audio-dir-to-splits",
            str(input_dir),
            str(output_root),
            "--source-dataset",
            self.name,
            "--eval-ratio",
            str(eval_ratio),
            "--query-duration",
            str(query_duration),
        ]
        if query_stride is not None:
            cmd.extend([
                "--query-stride",
                str(query_stride),
            ])
        cmd.extend([
            "--query-strategy",
            str(query_strategy),
            "--silence-top-db",
            str(silence_top_db),
        ])
        cmd.extend([
            "--seed",
            str(seed),
        ])
        result = subprocess.check_output(cmd, text=True)
        summary = json.loads(result)
        summary["input_dir"] = str(input_dir)
        summary["dataset"] = self.name
        return summary

    def inspect_local_audio(
        self,
        input_dir: Path,
        query_duration: float = 8.0,
        eval_ratio: float = 0.2,
    ) -> Dict:
        cmd = [
            "/usr/local/miniconda3/bin/python",
            "src/data/manifest_tools.py",
            "inspect-audio-dir",
            str(input_dir),
            "--query-duration",
            str(query_duration),
            "--eval-ratio",
            str(eval_ratio),
        ]
        result = subprocess.check_output(cmd, text=True)
        summary = json.loads(result)
        summary["dataset"] = self.name
        return summary

    def validate_local_manifests(self, manifests_dir: Path) -> Dict:
        cmd = [
            "/usr/local/miniconda3/bin/python",
            "src/data/manifest_tools.py",
            "validate-splits",
            str(manifests_dir),
        ]
        result = subprocess.check_output(cmd, text=True)
        summary = json.loads(result)
        summary["dataset"] = self.name
        return summary


class FMAAdapter(BaseAdapter):
    name = "fma"

    def describe(self) -> Dict:
        return {
            "name": "FMA",
            "source_url": "https://github.com/mdeff/fma",
            "recommended_subset": "fma_small",
            "catalog_strategy": "full tracks as references; random 5-15s crops as queries",
            "license_policy": "review per subset/track before commercial training",
        }


class MTGJamendoAdapter(BaseAdapter):
    name = "mtg_jamendo"

    def describe(self) -> Dict:
        return {
            "name": "MTG-Jamendo",
            "source_url": "https://github.com/MTG/mtg-jamendo-dataset",
            "recommended_subset": "small curated slice",
            "catalog_strategy": "download upstream audio subset then build catalog/query manifests",
            "license_policy": "verify CC terms for intended commercial use",
        }


class CCMusicAdapter(BaseAdapter):
    name = "ccmusic"

    def describe(self) -> Dict:
        return {
            "name": "CCMusic",
            "source_url": "https://ccmusic-database.github.io/en/database/ccm.html",
            "recommended_subset": "whitelisted approved subset only",
            "catalog_strategy": "use approved corpora only; normalize to project manifests",
            "license_policy": "application/permission review required before use",
        }


class ModelScopeMusicAdapter(BaseAdapter):
    name = "modelscope_music"

    def describe(self) -> Dict:
        return {
            "name": "ModelScope music datasets",
            "source_url": "https://modelscope.cn/search?page=1&search=music&type=dataset",
            "recommended_subset": "manual whitelist only",
            "catalog_strategy": "treat as discovery surface; add per-dataset adapter after legal review",
            "license_policy": "deny until whitelisted",
        }


ADAPTERS = {
    "fma": FMAAdapter(),
    "mtg_jamendo": MTGJamendoAdapter(),
    "ccmusic": CCMusicAdapter(),
    "modelscope_music": ModelScopeMusicAdapter(),
}

REGISTRY: List[DatasetRecord] = [
    DatasetRecord(
        name="FMA",
        source_url="https://github.com/mdeff/fma",
        license="Track-dependent / metadata CC BY 4.0; verify per subset",
        commercial_use="review_required",
        notes="Good first realistic MIR baseline",
    ),
    DatasetRecord(
        name="MTG-Jamendo",
        source_url="https://github.com/MTG/mtg-jamendo-dataset",
        license="Creative Commons source tracks; verify exact subset terms",
        commercial_use="review_required",
        notes="Good retrieval/tagging corpus with scripts",
    ),
    DatasetRecord(
        name="CCMusic",
        source_url="https://ccmusic-database.github.io/en/database/ccm.html",
        license="varies / application may be required",
        commercial_use="review_required",
        notes="Useful Chinese MIR source, needs permission review",
    ),
    DatasetRecord(
        name="ModelScope-music",
        source_url="https://modelscope.cn/search?page=1&search=music&type=dataset",
        license="varies by dataset",
        commercial_use="deny_until_whitelisted",
        notes="Discovery surface only until per-dataset review is complete",
    ),
]


def count_audio_files(input_dir: Path) -> int:
    return len([p for p in input_dir.rglob("*") if p.suffix.lower() in AUDIO_EXTS])


def assess_local_dataset_ready(
    dataset: str,
    input_dir: Path,
    query_duration: float = 8.0,
    eval_ratio: float = 0.2,
) -> Dict:
    adapter = ADAPTERS[dataset]
    input_dir = input_dir.resolve()
    exists = input_dir.exists()
    is_dir = input_dir.is_dir()
    inspect_summary = None
    num_audio_files = 0
    eligible_query_files = 0
    issues = []

    if not exists:
        issues.append("input_dir_missing")
    elif not is_dir:
        issues.append("input_path_not_directory")
    else:
        inspect_summary = adapter.inspect_local_audio(
            input_dir,
            query_duration=query_duration,
            eval_ratio=eval_ratio,
        )
        num_audio_files = int(inspect_summary.get("num_audio_files", 0))
        eligible_query_files = int(inspect_summary.get("eligible_query_files", 0))
        if num_audio_files < MIN_SMOKE_AUDIO_FILES:
            issues.append("not_enough_audio_files_for_smoke")
        if eligible_query_files < MIN_SMOKE_ELIGIBLE_QUERY_FILES:
            issues.append("not_enough_query_eligible_files_for_smoke")

    ready = len(issues) == 0
    recommendations = []
    if "input_dir_missing" in issues:
        recommendations.append(f"Create and populate {input_dir} with local audio files ({', '.join(AUDIO_EXTS)})")
    if "input_path_not_directory" in issues:
        recommendations.append("Replace the input path with a directory containing local audio files")
    if "not_enough_audio_files_for_smoke" in issues:
        recommendations.append(f"Add at least {MIN_SMOKE_AUDIO_FILES} audio files before running smoke-local")
    if "not_enough_query_eligible_files_for_smoke" in issues:
        recommendations.append(
            f"Add at least {MIN_SMOKE_ELIGIBLE_QUERY_FILES} files with duration >= {query_duration:.1f}s"
        )
    if ready:
        recommendations.append("Run smoke-local to verify the full train/index/eval/artifact pipeline on this local dataset")

    return {
        "dataset": dataset,
        "input_dir": str(input_dir),
        "exists": exists,
        "is_dir": is_dir,
        "ready_for_smoke": ready,
        "num_audio_files": num_audio_files,
        "eligible_query_files": eligible_query_files,
        "minimum_requirements": {
            "audio_files": MIN_SMOKE_AUDIO_FILES,
            "eligible_query_files": MIN_SMOKE_ELIGIBLE_QUERY_FILES,
            "query_duration": query_duration,
            "eval_ratio": eval_ratio,
        },
        "issues": issues,
        "recommendations": recommendations,
        "inspect": inspect_summary,
    }


def write_registry(output_path: str):
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "w") as f:
        json.dump([asdict(x) for x in REGISTRY], f, indent=2, ensure_ascii=False)
    return out


def inspect_batch(pairs: List[str], eval_ratio: float, query_duration: float) -> Dict:
    results = []
    for pair in pairs:
        dataset, input_dir = pair.split("=", 1)
        if dataset not in ADAPTERS:
            raise SystemExit(f"Unknown dataset adapter: {dataset}")
        summary = ADAPTERS[dataset].inspect_local_audio(
            Path(input_dir),
            eval_ratio=eval_ratio,
            query_duration=query_duration,
        )
        results.append(summary)
    return {"datasets": results, "count": len(results)}


def smoke_local_dataset(
    dataset: str,
    input_dir: Path,
    output_root: Path,
    eval_ratio: float,
    query_duration: float,
    query_stride: float | None,
    query_strategy: str,
    segment_strategy: str,
    silence_top_db: int,
    index_checkpoint_every_refs: int,
    max_test_queries: int | None,
    seed: int,
    train_epochs: int,
    batch_size: int,
    device: str,
) -> Dict:
    readiness = assess_local_dataset_ready(
        dataset,
        input_dir,
        query_duration=query_duration,
        eval_ratio=eval_ratio,
    )
    if not readiness["ready_for_smoke"]:
        raise SystemExit(json.dumps({
            "status": "blocked",
            "reason": "dataset_not_ready_for_smoke",
            "readiness": readiness,
        }, indent=2, ensure_ascii=False))

    adapter = ADAPTERS[dataset]
    resolved_device = resolve_device(device)
    inspect_summary = readiness["inspect"]
    prepare_summary = adapter.prepare_local_audio(
        input_dir,
        output_root / dataset,
        eval_ratio=eval_ratio,
        query_duration=query_duration,
        query_stride=query_stride,
        query_strategy=query_strategy,
        silence_top_db=silence_top_db,
        seed=seed,
    )
    manifests_dir = Path(prepare_summary["output_dir"])
    validate_summary = adapter.validate_local_manifests(manifests_dir)
    base_cfg = load_default_training_config()

    model_dir = output_root / f"{dataset}_models_smoke"
    index_dir = output_root / f"{dataset}_index_smoke"
    report_dir = output_root / f"{dataset}_reports_smoke"
    config_path = report_dir / "config.json"

    subprocess.run([
        "/usr/local/miniconda3/bin/python",
        "train.py",
        "--data", str(manifests_dir),
        "--output", str(model_dir),
        "--device", resolved_device,
        "--epochs", str(train_epochs),
        "--batch-size", str(batch_size),
        "--segment-strategy", str(segment_strategy),
        "--silence-top-db", str(silence_top_db),
    ], check=True)

    subprocess.run([
        "/usr/local/miniconda3/bin/python",
        "run_demo.py",
        "build-index",
        "--data", str(manifests_dir),
        "--model", str(model_dir / "best_model.pt"),
        "--output", str(index_dir),
        "--device", resolved_device,
        "--resume",
        "--checkpoint-every-refs", str(index_checkpoint_every_refs),
    ], check=True)

    report_dir.mkdir(parents=True, exist_ok=True)
    eval_json = report_dir / "eval.json"
    subprocess.run([
        "/usr/local/miniconda3/bin/python",
        "evaluate.py",
        "--data", str(manifests_dir),
        "--model", str(model_dir / "best_model.pt"),
        "--index-prefix", str(index_dir / "reference"),
        "--split", "test",
        "--device", resolved_device,
        "--fast-eval",
        "--output-json", str(eval_json),
        "--seed", str(seed),
        *([] if max_test_queries is None else ["--max-queries", str(max_test_queries)]),
    ], check=True)

    config = build_smoke_config_summary(
        dataset=dataset,
        manifests_dir=manifests_dir,
        manifest_query_duration=query_duration,
        train_epochs=train_epochs,
        batch_size=batch_size,
        requested_device=device,
        resolved_device=resolved_device,
        base_cfg=base_cfg,
    )
    config["data"]["manifest_query_stride"] = query_stride
    config["data"]["manifest_query_strategy"] = query_strategy
    config["data"]["silence_top_db"] = silence_top_db
    config["run"]["index_checkpoint_every_refs"] = index_checkpoint_every_refs
    config["run"]["index_resume_enabled"] = True
    config["run"]["train_segment_strategy"] = segment_strategy
    config["run"]["max_test_queries"] = max_test_queries
    report_dir.mkdir(parents=True, exist_ok=True)
    config_path.write_text(json.dumps(config, indent=2))

    subprocess.run([
        "/usr/local/miniconda3/bin/python",
        "scripts/generate_artifacts.py",
        "--eval-json", str(eval_json),
        "--config-json", str(config_path),
        "--output-dir", str(report_dir),
        "--model-version", f"{dataset}-smoke",
        "--data-version", f"{dataset}_local",
    ], check=True)

    return {
        "dataset": dataset,
        "readiness": readiness,
        "inspect": inspect_summary,
        "prepare": prepare_summary,
        "validate": validate_summary,
        "requested_device": device,
        "resolved_device": resolved_device,
        "model_dir": str(model_dir),
        "index_dir": str(index_dir),
        "report_dir": str(report_dir),
        "eval_json": str(eval_json),
    }


def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("registry")
    p.add_argument("--output", default="data/dataset_registry.json")

    p = sub.add_parser("init")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("--root", default="data/external")

    p = sub.add_parser("describe")
    p.add_argument("dataset", choices=sorted(ADAPTERS))

    p = sub.add_parser("prepare-local")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("input_dir")
    p.add_argument("--output-root", default="data/external_ingested")
    p.add_argument("--eval-ratio", type=float, default=0.2)
    p.add_argument("--query-duration", type=float, default=8.0)
    p.add_argument("--query-stride", type=float, default=None)
    p.add_argument("--query-strategy", choices=["random", "sliding", "silence_aware", "high_energy", "onset_aware", "beat_aware", "repeated_section_aware", "hybrid"], default="random")
    p.add_argument("--silence-top-db", type=int, default=30)
    p.add_argument("--seed", type=int, default=42)

    p = sub.add_parser("inspect-local")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("input_dir")
    p.add_argument("--eval-ratio", type=float, default=0.2)
    p.add_argument("--query-duration", type=float, default=8.0)

    p = sub.add_parser("inspect-batch")
    p.add_argument("pairs", nargs="+", help="dataset=input_dir")
    p.add_argument("--eval-ratio", type=float, default=0.2)
    p.add_argument("--query-duration", type=float, default=8.0)

    p = sub.add_parser("validate-local")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("manifests_dir")

    p = sub.add_parser("check-local-ready")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("input_dir")
    p.add_argument("--eval-ratio", type=float, default=0.2)
    p.add_argument("--query-duration", type=float, default=8.0)

    p = sub.add_parser("smoke-local")
    p.add_argument("dataset", choices=sorted(ADAPTERS))
    p.add_argument("input_dir")
    p.add_argument("--output-root", default="data/external_smoke")
    p.add_argument("--eval-ratio", type=float, default=0.2)
    p.add_argument("--query-duration", type=float, default=8.0)
    p.add_argument("--query-stride", type=float, default=None)
    p.add_argument("--query-strategy", choices=["random", "sliding", "silence_aware", "high_energy", "onset_aware", "beat_aware", "repeated_section_aware", "hybrid"], default="random")
    p.add_argument("--segment-strategy", choices=["random", "silence_aware", "high_energy", "onset_aware", "beat_aware", "repeated_section_aware", "hybrid"], default="random")
    p.add_argument("--silence-top-db", type=int, default=30)
    p.add_argument("--index-checkpoint-every-refs", type=int, default=100)
    p.add_argument("--max-test-queries", type=int, default=None)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--train-epochs", type=int, default=1)
    p.add_argument("--batch-size", type=int, default=2)
    p.add_argument("--device", default="cpu")

    args = parser.parse_args()
    if args.cmd == "registry":
        path = write_registry(args.output)
        print(path)
    elif args.cmd == "init":
        root = Path(args.root) / args.dataset
        print(json.dumps(ADAPTERS[args.dataset].init_layout(root), indent=2, ensure_ascii=False))
    elif args.cmd == "describe":
        print(json.dumps(ADAPTERS[args.dataset].describe(), indent=2, ensure_ascii=False))
    elif args.cmd == "prepare-local":
        root = Path(args.output_root) / args.dataset
        summary = ADAPTERS[args.dataset].prepare_local_audio(
            Path(args.input_dir),
            root,
            eval_ratio=args.eval_ratio,
            query_duration=args.query_duration,
            query_stride=args.query_stride,
            query_strategy=args.query_strategy,
            silence_top_db=args.silence_top_db,
            seed=args.seed,
        )
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    elif args.cmd == "inspect-local":
        summary = ADAPTERS[args.dataset].inspect_local_audio(
            Path(args.input_dir),
            eval_ratio=args.eval_ratio,
            query_duration=args.query_duration,
        )
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    elif args.cmd == "inspect-batch":
        summary = inspect_batch(args.pairs, args.eval_ratio, args.query_duration)
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    elif args.cmd == "validate-local":
        summary = ADAPTERS[args.dataset].validate_local_manifests(Path(args.manifests_dir))
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    elif args.cmd == "check-local-ready":
        summary = assess_local_dataset_ready(
            dataset=args.dataset,
            input_dir=Path(args.input_dir),
            eval_ratio=args.eval_ratio,
            query_duration=args.query_duration,
        )
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    elif args.cmd == "smoke-local":
        summary = smoke_local_dataset(
            dataset=args.dataset,
            input_dir=Path(args.input_dir),
            output_root=Path(args.output_root),
            eval_ratio=args.eval_ratio,
            query_duration=args.query_duration,
            query_stride=args.query_stride,
            query_strategy=args.query_strategy,
            segment_strategy=args.segment_strategy,
            silence_top_db=args.silence_top_db,
            index_checkpoint_every_refs=args.index_checkpoint_every_refs,
            max_test_queries=args.max_test_queries,
            seed=args.seed,
            train_epochs=args.train_epochs,
            batch_size=args.batch_size,
            device=args.device,
        )
        print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()