from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from . import __version__
from .api import cache_embeddings, evaluate, run, train
from .assets import default_assets_root, discover_cached_embeddings, discover_checkpoints, discover_repo_checkpoints
from .hub import DEFAULT_REPO_ID, download_checkpoint


def _add_dataset_args(parser: argparse.ArgumentParser) -> None:
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--dataset", help="BEIR dataset name, such as scifact.")
    group.add_argument("--dataset-path", type=Path, help="Local BEIR-format dataset directory.")
    parser.add_argument("--datasets-dir", type=Path, help="BEIR download root (default: <assets>/datasets).")


def _add_encoder_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--encoder", help="Built-in encoder alias: minilm, e5, or mpnet.")
    parser.add_argument("--encoder-model-name", help="Custom SentenceTransformers model name.")
    parser.add_argument("--embedding-dim", type=int, help="Required with --encoder-model-name.")
    parser.add_argument("--query-prefix", default="")
    parser.add_argument("--passage-prefix", default="")
    parser.add_argument("--encoder-revision", help="Optional immutable revision for a custom base encoder.")


def _add_runtime_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--assets-root", type=Path, default=default_assets_root())
    parser.add_argument("--cache-dir", type=Path)
    parser.add_argument("--max-queries", type=int)
    parser.add_argument("--seed", type=int, default=42)


def _dataset_values(args: argparse.Namespace) -> tuple[str, Path]:
    if args.dataset_path:
        path = args.dataset_path.resolve()
        return path.name, path.parent
    return args.dataset, args.datasets_dir or (args.assets_root / "datasets")


def _encoder_label(args: argparse.Namespace) -> str:
    return args.encoder or args.encoder_model_name or "custom"


def _cache_dir(args: argparse.Namespace, dataset: str) -> Path:
    return args.cache_dir or args.assets_root / f"cache_{_encoder_label(args).replace('/', '-')}_{dataset}"


def _command_info(args: argparse.Namespace) -> int:
    print(
        json.dumps(
            {
                "version": __version__,
                "default_hugging_face_repo": DEFAULT_REPO_ID,
                "available_checkpoint": "minilm/scifact",
                "training_recipe": {
                    "objective": "improvement-margin",
                    "improvement_margin": 0.05,
                    "hard_negatives": 63,
                    "hard_negative_method": "dense",
                    "hard_negative_pool": 100,
                    "projector_initialization": "identity",
                    "optimizer": "adam",
                    "learning_rate": 1e-4,
                    "weight_decay": 1e-5,
                    "batch_size": 32,
                    "epochs": 30,
                    "patience": 7,
                    "seed": 42,
                },
                "supported_encoders": ["minilm", "e5", "mpnet"],
            },
            indent=2,
        )
    )
    return 0


def _command_list_assets(args: argparse.Namespace) -> int:
    print(
        json.dumps(
            {
                "assets_root": str(args.assets_root),
                "cached_embeddings": [
                    item.__dict__ | {"path": str(item.path)} for item in discover_cached_embeddings(args.assets_root)
                ],
                "workspace_checkpoints": [
                    item.__dict__ | {"path": str(item.path)} for item in discover_checkpoints(args.assets_root)
                ],
                "repo_checkpoints": [
                    item.__dict__ | {"path": str(item.path)}
                    for item in discover_repo_checkpoints(Path(__file__).resolve().parents[2])
                ],
            },
            indent=2,
            default=str,
        )
    )
    return 0


def _command_download(args: argparse.Namespace) -> int:
    downloaded = download_checkpoint(
        encoder=args.encoder,
        dataset=args.dataset,
        repo_id=args.repo_id,
        revision=args.revision,
        cache_dir=args.hub_cache_dir,
    )
    print(json.dumps({"checkpoint": str(downloaded.checkpoint_path), "repo_id": downloaded.repo_id}, indent=2))
    return 0


def _command_cache(args: argparse.Namespace) -> int:
    dataset, datasets_dir = _dataset_values(args)
    built = cache_embeddings(
        encoder=args.encoder,
        dataset=dataset,
        cache_dir=_cache_dir(args, dataset),
        datasets_dir=datasets_dir,
        device=args.device,
        encoder_model_name=args.encoder_model_name,
        embedding_dim=args.embedding_dim,
        query_prefix=args.query_prefix,
        passage_prefix=args.passage_prefix,
        encoder_revision=args.encoder_revision,
        batch_size=args.batch_size,
        num_negatives=args.num_negatives,
        negative_pool=args.negative_pool,
        max_queries=args.max_queries,
        seed=args.seed,
    )
    print(json.dumps({"cache_dir": str(built), "dataset": dataset}, indent=2))
    return 0


def _training_kwargs(args: argparse.Namespace) -> dict:
    dataset, datasets_dir = _dataset_values(args)
    return {
        "encoder": args.encoder,
        "dataset": dataset,
        "cache_dir": _cache_dir(args, dataset),
        "datasets_dir": datasets_dir,
        "output_dir": args.output_dir or args.assets_root,
        "device": args.device,
        "encoder_model_name": args.encoder_model_name,
        "embedding_dim": args.embedding_dim,
        "query_prefix": args.query_prefix,
        "passage_prefix": args.passage_prefix,
        "encoder_revision": args.encoder_revision,
        "max_queries": args.max_queries,
        "batch_size": args.batch_size,
        "epochs": args.epochs,
        "lr": args.lr,
        "weight_decay": args.weight_decay,
        "num_negatives": args.num_negatives,
        "feedback_k": args.candidate_k,
        "ranking_k": args.ranking_k,
        "k": args.k,
        "improvement_margin": args.improvement_margin,
        "patience": args.patience,
        "seed": args.seed,
    }


def _print_result(payload: dict) -> int:
    print(json.dumps(payload, indent=2, default=str))
    return 0


def _command_train(args: argparse.Namespace) -> int:
    return _print_result(train(**_training_kwargs(args)))


def _command_run(args: argparse.Namespace) -> int:
    kwargs = _training_kwargs(args)
    kwargs["negative_pool"] = args.negative_pool
    return _print_result(run(**kwargs))


def _command_evaluate(args: argparse.Namespace) -> int:
    dataset, datasets_dir = _dataset_values(args)
    return _print_result(
        evaluate(
            encoder=args.encoder,
            dataset=dataset,
            cache_dir=_cache_dir(args, dataset),
            datasets_dir=datasets_dir,
            checkpoint_path=args.checkpoint,
            device=args.device,
            encoder_model_name=args.encoder_model_name,
            embedding_dim=args.embedding_dim,
            query_prefix=args.query_prefix,
            passage_prefix=args.passage_prefix,
            encoder_revision=args.encoder_revision,
            max_queries=args.max_queries,
            seed=args.seed,
            feedback_k=args.candidate_k,
            ranking_k=args.ranking_k,
            k=args.k,
        )
    )


def _configure_training_parser(parser: argparse.ArgumentParser) -> None:
    _add_dataset_args(parser)
    _add_encoder_args(parser)
    _add_runtime_args(parser)
    parser.add_argument("--output-dir", type=Path)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--epochs", type=int, default=30)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-5)
    parser.add_argument("--num-negatives", type=int, default=63)
    parser.add_argument("--candidate-k", type=int, default=100)
    parser.add_argument("--ranking-k", type=int, default=10)
    parser.add_argument("--k", type=int, default=10)
    parser.add_argument("--improvement-margin", type=float, default=0.05)
    parser.add_argument("--patience", type=int, default=7)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="imrnns",
        description="Train, evaluate, and use IMRNN dense-retrieval adapters.",
    )
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    subparsers = parser.add_subparsers(dest="command", required=True)

    info = subparsers.add_parser("info", help="Show package and recipe information.")
    info.set_defaults(func=_command_info)

    assets = subparsers.add_parser("list-assets", help="List local caches and checkpoints.")
    assets.add_argument("--assets-root", type=Path, default=default_assets_root())
    assets.set_defaults(func=_command_list_assets)

    download = subparsers.add_parser("download", help="Download the validated adapter checkpoint.")
    download.add_argument("--encoder", required=True, choices=["minilm"])
    download.add_argument("--dataset", required=True, choices=["scifact"])
    download.add_argument("--repo-id", default=DEFAULT_REPO_ID)
    download.add_argument("--revision")
    download.add_argument("--hub-cache-dir", type=Path)
    download.set_defaults(func=_command_download)

    cache = subparsers.add_parser("cache", help="Encode a dataset and mine dense hard negatives.")
    _add_dataset_args(cache)
    _add_encoder_args(cache)
    _add_runtime_args(cache)
    cache.add_argument("--batch-size", type=int, default=64)
    cache.add_argument("--num-negatives", type=int, default=63)
    cache.add_argument("--negative-pool", type=int, default=100)
    cache.set_defaults(func=_command_cache)

    train_parser = subparsers.add_parser("train", help="Train an adapter from a prepared cache.")
    _configure_training_parser(train_parser)
    train_parser.set_defaults(func=_command_train)

    evaluate_parser = subparsers.add_parser("evaluate", help="Evaluate a checkpoint against its base retriever.")
    _add_dataset_args(evaluate_parser)
    _add_encoder_args(evaluate_parser)
    _add_runtime_args(evaluate_parser)
    evaluate_parser.add_argument("--checkpoint", type=Path, required=True)
    evaluate_parser.add_argument("--candidate-k", type=int, default=100)
    evaluate_parser.add_argument("--ranking-k", type=int, default=10)
    evaluate_parser.add_argument("--k", type=int, default=10)
    evaluate_parser.set_defaults(func=_command_evaluate)

    run_parser = subparsers.add_parser("run", help="Cache, train, and evaluate end to end.")
    _configure_training_parser(run_parser)
    run_parser.add_argument("--negative-pool", type=int, default=100)
    run_parser.set_defaults(func=_command_run)
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return int(args.func(args))
    except (FileNotFoundError, ImportError, RuntimeError, ValueError) as exc:
        print(f"imrnns: error: {exc}", file=sys.stderr)
        return 2