import json
import re
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple


def venv_python_path(venv_dir: Path) -> Path:
    if (venv_dir / "Scripts" / "python.exe").exists():
        return venv_dir / "Scripts" / "python.exe"
    return venv_dir / "bin" / "python"


def _sanitize_token(value: str) -> str:
    normalized = re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._-")
    return normalized or "input"


def create_embeddings_run_dir(workspace_root: Path, model_name: str, input_file: Path) -> Tuple[Path, str]:
    base_dir = workspace_root / "data" / "runs" / "embeddings" / model_name.upper()
    base_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    token = _sanitize_token(input_file.stem)
    run_id = f"{timestamp}_{token}"
    run_dir = base_dir / run_id

    suffix = 2
    while run_dir.exists():
        run_id = f"{timestamp}_{token}_{suffix}"
        run_dir = base_dir / run_id
        suffix += 1

    run_dir.mkdir(parents=True, exist_ok=False)
    return run_dir, run_id


def _load_model_metadata(metadata_root: Path, model_name: str) -> Dict[str, Any]:
    metadata_dir = metadata_root / model_name
    candidates = [metadata_dir / "model.json", metadata_dir / f"{model_name}.json"]
    for path in candidates:
        if path.exists():
            with path.open("r", encoding="utf-8") as handle:
                return json.load(handle)
    raise FileNotFoundError(f"Metadata not found for model {model_name} in {metadata_dir}")


def build_embeddings_command(
    workspace_root: Path,
    model_name: str,
    input_file: Path,
    input_format: str,
    output_file: Optional[Path] = None,
    extra_args: Optional[List[str]] = None,
    metadata_root: Optional[Path] = None,
) -> Tuple[str, List[str], str, Path, Path, str]:
    """Build a venv-scoped embeddings command for a model.

    Returns (program, args, cwd, resolved_output_path, run_dir, run_id).
    """
    model_dir = workspace_root / "models" / model_name
    embeddings_script = model_dir / "embeddings.py"
    venv_dir = model_dir / ".venv"
    python_exe = venv_python_path(venv_dir)

    if not embeddings_script.exists():
        raise FileNotFoundError(f"Embeddings script not found: {embeddings_script}")
    if not python_exe.exists():
        raise FileNotFoundError(f"Model venv python not found: {python_exe}")
    if not input_file.exists():
        raise FileNotFoundError(f"Input file not found: {input_file}")

    if output_file is None:
        run_dir, run_id = create_embeddings_run_dir(workspace_root, model_name, input_file)
        resolved_output = run_dir / f"{input_file.stem}.npy"
    else:
        resolved_output = output_file
        run_dir = resolved_output.parent
        run_id = run_dir.name
        run_dir.mkdir(parents=True, exist_ok=True)

    resolved_output.parent.mkdir(parents=True, exist_ok=True)

    resolved_metadata_root = metadata_root or (workspace_root / "metadata")
    metadata = _load_model_metadata(resolved_metadata_root, model_name)
    cli_version = int(metadata.get("embeddingsCliVersion", 1))
    if model_name.upper() == "CLAMP" and cli_version < 2:
        raise RuntimeError(
            "CLAMP requires embeddingsCliVersion >= 2 and the new embeddings CLI protocol."
        )

    if cli_version >= 2:
        model_input_type = str(metadata.get("inputType", "smiles")).strip().lower()
        normalized_input_format = str(input_format).strip().lower()
        args = [
            str(embeddings_script),
            "--file",
            str(input_file),
            "--input-format",
            normalized_input_format,
            "--model-input-type",
            model_input_type,
            "--output-dir",
            str(resolved_output.parent),
            "--output-stem",
            resolved_output.stem,
        ]
    else:
        # Legacy embeddings CLI compatibility.
        args = [
            str(embeddings_script),
            "--input",
            str(input_file),
            "--output",
            str(resolved_output),
        ]

    if extra_args:
        args.extend(extra_args)

    return str(python_exe), args, str(model_dir), resolved_output, run_dir, run_id
