from __future__ import annotations

import hashlib
import json
import subprocess
from importlib import metadata as package_metadata
from pathlib import Path
from typing import Any, Optional

import torch

from .beir_data import load_beir_splits
from .caching import build_cache
from .checkpoints import default_checkpoint_name, load_model, save_checkpoint
from .data import ContrastiveCachedDataset, load_cached_split
from .encoders import encoder_storage_key, resolve_encoder_spec
from .evaluation import evaluate_model, evaluate_model_with_baseline
from .model import IMRNN, ModelConfig
from .training import TrainingConfig, initialize_projector, set_seed, train_model


def _dependency_version(distribution: str) -> str | None:
    try:
        return package_metadata.version(distribution)
    except package_metadata.PackageNotFoundError:
        return None


def _git_commit() -> str | None:
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=Path(__file__).resolve().parents[2],
            capture_output=True,
            text=True,
            check=True,
            timeout=2,
        )
    except (OSError, subprocess.SubprocessError):
        return None
    return result.stdout.strip() or None


def _load_cache_manifest(cache_dir: Path) -> tuple[dict[str, Any] | None, str | None]:
    path = cache_dir / "manifest.json"
    if not path.is_file():
        return None, None
    raw = path.read_bytes()
    manifest = json.loads(raw)
    if not isinstance(manifest, dict):
        raise ValueError(f"Cache manifest must contain a JSON object: {path}")
    return manifest, hashlib.sha256(raw).hexdigest()


def _validate_cache(
    cache_dir: Path,
    *,
    dataset: str,
    model_name: str,
    model_revision: str | None,
    max_queries: int | None,
    seed: int,
    num_negatives: int | None = None,
) -> dict[str, Any]:
    manifest, _ = _load_cache_manifest(cache_dir)
    if manifest is None:
        raise ValueError(f"Cache manifest is missing: {cache_dir / 'manifest.json'}")
    expected = {
        "schema_version": 2,
        "dataset": dataset,
        "model_name": model_name,
        "model_revision": model_revision,
        "negative_method": "dense",
        "max_queries": max_queries,
        "seed": seed,
    }
    mismatches = {
        key: {"expected": value, "found": manifest.get(key)}
        for key, value in expected.items()
        if manifest.get(key) != value
    }
    if num_negatives is not None:
        cached_negatives = int(manifest.get("num_negatives", 0))
        if cached_negatives < num_negatives:
            mismatches["num_negatives"] = {"expected_at_least": num_negatives, "found": cached_negatives}
    if mismatches:
        raise ValueError(f"Training cache is incompatible with this run: {mismatches}")
    return manifest


def cache_embeddings(
    *,
    encoder: Optional[str],
    dataset: str,
    cache_dir: Path,
    datasets_dir: Path,
    device: str = "cpu",
    encoder_model_name: Optional[str] = None,
    embedding_dim: Optional[int] = None,
    query_prefix: str = "",
    passage_prefix: str = "",
    encoder_revision: str | None = None,
    batch_size: int = 64,
    num_negatives: int = 63,
    negative_pool: int = 100,
    max_queries: Optional[int] = None,
    seed: int = 42,
) -> Path:
    encoder_spec = resolve_encoder_spec(
        encoder=encoder,
        encoder_model_name=encoder_model_name,
        embedding_dim=embedding_dim,
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        encoder_revision=encoder_revision,
    )
    return build_cache(
        dataset_name=dataset,
        encoder_spec=encoder_spec,
        cache_dir=cache_dir,
        datasets_dir=datasets_dir,
        device=device,
        batch_size=batch_size,
        num_negatives=num_negatives,
        negative_pool=negative_pool,
        max_queries=max_queries,
        seed=seed,
    )


def train(
    *,
    encoder: Optional[str],
    dataset: str,
    cache_dir: Path,
    datasets_dir: Path,
    output_dir: Path,
    device: str = "cpu",
    encoder_model_name: Optional[str] = None,
    embedding_dim: Optional[int] = None,
    query_prefix: str = "",
    passage_prefix: str = "",
    encoder_revision: str | None = None,
    max_queries: Optional[int] = None,
    batch_size: int = 32,
    epochs: int = 30,
    lr: float = 1e-4,
    weight_decay: float = 1e-5,
    num_negatives: int = 63,
    feedback_k: int = 100,
    ranking_k: int = 10,
    k: int = 10,
    improvement_margin: float = 0.05,
    patience: int = 7,
    seed: int = 42,
) -> dict[str, Any]:
    encoder_spec = resolve_encoder_spec(
        encoder=encoder,
        encoder_model_name=encoder_model_name,
        embedding_dim=embedding_dim,
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        encoder_revision=encoder_revision,
    )
    _validate_cache(
        cache_dir,
        dataset=dataset,
        model_name=encoder_spec.model_name,
        model_revision=encoder_spec.revision,
        max_queries=max_queries,
        seed=seed,
        num_negatives=num_negatives,
    )
    sources = load_beir_splits(dataset, datasets_dir=datasets_dir, max_queries=max_queries, seed=seed)
    train_split = load_cached_split(cache_dir, "train", sources["train"], encoder_spec, device)
    val_split = load_cached_split(cache_dir, "validation", sources["validation"], encoder_spec, device)
    test_split = load_cached_split(cache_dir, "test", sources["test"], encoder_spec, device)

    set_seed(seed)
    model = IMRNN(ModelConfig(input_dim=encoder_spec.embedding_dim))
    initialize_projector(model)

    train_dataset = ContrastiveCachedDataset(train_split, num_negatives)
    val_dataset = ContrastiveCachedDataset(val_split, num_negatives)
    if len(train_dataset) == 0:
        raise ValueError("No training examples were constructed from the cached training split.")
    if len(val_dataset) == 0:
        raise ValueError("No validation examples were constructed from the cached validation split.")

    training_metrics = train_model(
        model=model,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        config=TrainingConfig(
            batch_size=batch_size,
            epochs=epochs,
            lr=lr,
            weight_decay=weight_decay,
            num_negatives=num_negatives,
            improvement_margin=improvement_margin,
            patience=patience,
            seed=seed,
        ),
        device=device,
        validation_metric_fn=lambda current_model: (
            sum(
                evaluate_model(
                    model=current_model,
                    cached_split=val_split,
                    device=device,
                    feedback_k=feedback_k,
                    ranking_k=ranking_k,
                    k_values=[k],
                ).values()
            )
            / 3.0
        ),
        validation_metric_name=f"mean(NDCG@{k},Recall@{k},MRR@{k})",
    )
    evaluation_metrics, base_metrics = evaluate_model_with_baseline(
        model=model,
        cached_split=test_split,
        device=device,
        feedback_k=feedback_k,
        ranking_k=ranking_k,
        k_values=[k],
    )
    evaluation_delta = {name: evaluation_metrics[name] - base_metrics[name] for name in evaluation_metrics}
    beats_base = bool(evaluation_delta) and all(delta > 0 for delta in evaluation_delta.values())

    checkpoint_stem = encoder_storage_key(encoder or encoder_spec.key)
    checkpoint_path = output_dir / default_checkpoint_name(checkpoint_stem, dataset)
    cache_manifest, cache_manifest_sha256 = _load_cache_manifest(cache_dir)
    metadata = {
        "encoder": checkpoint_stem,
        "encoder_model_name": encoder_spec.model_name,
        "encoder_revision": encoder_spec.revision,
        "query_prefix": encoder_spec.query_prefix,
        "passage_prefix": encoder_spec.passage_prefix,
        "dataset": dataset,
        "cache_dir": str(cache_dir),
        "model_config": {
            "input_dim": encoder_spec.embedding_dim,
            "output_dim": model.config.output_dim,
            "hidden_dim": model.config.hidden_dim,
            "dropout": model.config.dropout,
        },
        "training": training_metrics,
        "evaluation": evaluation_metrics,
        "base_evaluation": base_metrics,
        "evaluation_delta": evaluation_delta,
        "beats_base_all_metrics": beats_base,
        "seed": seed,
        "objective": "improvement-margin",
        "optimizer": "adam",
        "learning_rate": lr,
        "weight_decay": weight_decay,
        "batch_size": batch_size,
        "epochs": epochs,
        "patience": patience,
        "improvement_margin": improvement_margin,
        "projector_initialization": "identity",
        "number_of_negatives": num_negatives,
        "negative_retrieval_method": cache_manifest.get("negative_method") if cache_manifest else None,
        "negative_pool": cache_manifest.get("negative_pool") if cache_manifest else None,
        "candidate_k": feedback_k,
        "feedback_k": feedback_k,
        "ranking_k": ranking_k,
        "base_encoder": encoder_spec.model_name,
        "base_encoder_revision": encoder_spec.revision,
        "dataset_manifest": cache_manifest,
        "dataset_manifest_sha256": cache_manifest_sha256,
        "package_version": _dependency_version("imrnns"),
        "git_commit": _git_commit(),
        "environment": {
            "torch": str(torch.__version__),
            "sentence-transformers": _dependency_version("sentence-transformers"),
            "transformers": _dependency_version("transformers"),
        },
    }
    save_checkpoint(checkpoint_path, model, metadata)
    return {
        "checkpoint": checkpoint_path,
        "training": training_metrics,
        "evaluation": evaluation_metrics,
        "base_evaluation": base_metrics,
        "evaluation_delta": evaluation_delta,
        "beats_base_all_metrics": beats_base,
        "metadata": metadata,
    }


def evaluate(
    *,
    encoder: Optional[str],
    dataset: str,
    cache_dir: Path,
    datasets_dir: Path,
    checkpoint_path: Path,
    device: str = "cpu",
    encoder_model_name: Optional[str] = None,
    embedding_dim: Optional[int] = None,
    query_prefix: str = "",
    passage_prefix: str = "",
    encoder_revision: str | None = None,
    max_queries: Optional[int] = None,
    seed: int = 42,
    feedback_k: int = 100,
    ranking_k: int = 10,
    k: int = 10,
) -> dict[str, Any]:
    encoder_spec = resolve_encoder_spec(
        encoder=encoder,
        encoder_model_name=encoder_model_name,
        embedding_dim=embedding_dim,
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        encoder_revision=encoder_revision,
    )
    _validate_cache(
        cache_dir,
        dataset=dataset,
        model_name=encoder_spec.model_name,
        model_revision=encoder_spec.revision,
        max_queries=max_queries,
        seed=seed,
    )
    sources = load_beir_splits(dataset, datasets_dir=datasets_dir, max_queries=max_queries, seed=seed)
    test_split = load_cached_split(cache_dir, "test", sources["test"], encoder_spec, device)
    model, metadata, missing, unexpected = load_model(
        checkpoint_path=checkpoint_path,
        device=device,
    )
    if model.config.input_dim != encoder_spec.embedding_dim:
        raise ValueError(
            f"Checkpoint expects {model.config.input_dim}-dimensional embeddings; encoder produces "
            f"{encoder_spec.embedding_dim}."
        )
    metrics, base_metrics = evaluate_model_with_baseline(
        model=model,
        cached_split=test_split,
        device=device,
        feedback_k=feedback_k,
        ranking_k=ranking_k,
        k_values=[k],
    )
    metric_delta = {name: metrics[name] - base_metrics[name] for name in metrics}
    return {
        "checkpoint": checkpoint_path,
        "metrics": metrics,
        "base_metrics": base_metrics,
        "metric_delta": metric_delta,
        "beats_base_all_metrics": bool(metric_delta) and all(delta > 0 for delta in metric_delta.values()),
        "metadata": metadata,
        "missing_keys": missing,
        "unexpected_keys": unexpected,
    }


def run(
    *,
    encoder: Optional[str],
    dataset: str,
    cache_dir: Path,
    datasets_dir: Path,
    output_dir: Path,
    device: str = "cpu",
    encoder_model_name: Optional[str] = None,
    embedding_dim: Optional[int] = None,
    query_prefix: str = "",
    passage_prefix: str = "",
    encoder_revision: str | None = None,
    max_queries: Optional[int] = None,
    batch_size: int = 32,
    epochs: int = 30,
    lr: float = 1e-4,
    weight_decay: float = 1e-5,
    num_negatives: int = 63,
    negative_pool: int = 100,
    feedback_k: int = 100,
    ranking_k: int = 10,
    k: int = 10,
    improvement_margin: float = 0.05,
    patience: int = 7,
    seed: int = 42,
) -> dict[str, Any]:
    cache_embeddings(
        encoder=encoder,
        dataset=dataset,
        cache_dir=cache_dir,
        datasets_dir=datasets_dir,
        device=device,
        encoder_model_name=encoder_model_name,
        embedding_dim=embedding_dim,
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        encoder_revision=encoder_revision,
        batch_size=batch_size,
        num_negatives=num_negatives,
        negative_pool=negative_pool,
        max_queries=max_queries,
        seed=seed,
    )
    return train(
        encoder=encoder,
        dataset=dataset,
        cache_dir=cache_dir,
        datasets_dir=datasets_dir,
        output_dir=output_dir,
        device=device,
        encoder_model_name=encoder_model_name,
        embedding_dim=embedding_dim,
        query_prefix=query_prefix,
        passage_prefix=passage_prefix,
        encoder_revision=encoder_revision,
        max_queries=max_queries,
        batch_size=batch_size,
        epochs=epochs,
        lr=lr,
        weight_decay=weight_decay,
        num_negatives=num_negatives,
        feedback_k=feedback_k,
        ranking_k=ranking_k,
        k=k,
        improvement_margin=improvement_margin,
        patience=patience,
        seed=seed,
    )