from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import yaml

from common import ROOT, localize_workspace_path

DEFAULT_CONFIG = ROOT / "plugins" / "1c" / "training" / "configs" / "qwen3-coder-30b-a3b-lora.yaml"
DEFAULT_REQUIRED_BASE_FILES = [
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
]


def load_config(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a YAML mapping")
    return data


def count_jsonl(path: Path) -> int:
    if not path.exists():
        return 0
    count = 0
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                json.loads(line)
                count += 1
    return count


def check_python_deps() -> list[str]:
    missing = []
    for module in ("torch", "transformers", "datasets", "peft", "accelerate"):
        try:
            __import__(module)
        except Exception as exc:
            missing.append(f"{module}: {type(exc).__name__}: {exc}")
    return missing


def check_cuda() -> tuple[bool, str]:
    try:
        import torch
    except Exception as exc:
        return False, f"torch unavailable: {exc}"
    if not torch.cuda.is_available():
        return False, "torch.cuda.is_available() is false"
    return True, torch.cuda.get_device_name(0)


def expected_base_files(config: dict) -> list[str]:
    configured = config.get("required_base_files")
    if isinstance(configured, list):
        return [str(item) for item in configured if str(item).strip()]
    return list(DEFAULT_REQUIRED_BASE_FILES)


def missing_weight_shards(base_model_path: Path) -> list[str]:
    index_path = base_model_path / "model.safetensors.index.json"
    if not index_path.exists():
        return []
    try:
        index = json.loads(index_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        return [f"invalid model.safetensors.index.json: {exc}"]
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        return ["model.safetensors.index.json has no weight_map"]
    shards = sorted({str(value) for value in weight_map.values() if str(value).strip()})
    return [name for name in shards if not (base_model_path / name).exists()]


def main() -> int:
    parser = argparse.ArgumentParser(description="Preflight checks for 1C LoRA training.")
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    args = parser.parse_args()

    errors: list[str] = []
    warnings: list[str] = []
    config = load_config(args.config)

    dataset_path = localize_workspace_path(config["dataset_path"])
    dataset_count = count_jsonl(dataset_path)
    if dataset_count == 0:
        errors.append(f"training dataset is missing or empty: {dataset_path}")
    elif dataset_count < 50:
        warnings.append(f"dataset has only {dataset_count} record(s); this is not enough for useful fine-tuning")

    base_model_path = localize_workspace_path(config["base_model_path"])
    missing_files = [name for name in expected_base_files(config) if not (base_model_path / name).exists()]
    if missing_files:
        errors.append(f"base model is incomplete at {base_model_path}: missing {', '.join(missing_files)}")
    missing_shards = missing_weight_shards(base_model_path)
    if missing_shards:
        errors.append(f"base model shard set is incomplete at {base_model_path}: missing {', '.join(missing_shards)}")

    output_dir = localize_workspace_path(config["output_dir"])
    if not output_dir.parent.exists():
        warnings.append(f"adapter parent directory does not exist yet: {output_dir.parent}")

    missing_deps = check_python_deps()
    if missing_deps:
        errors.append(f"missing Python training dependencies: {'; '.join(missing_deps)}")

    cuda_ok, cuda_message = check_cuda()
    if cuda_ok:
        print(f"CUDA: {cuda_message}")
    else:
        errors.append(f"GPU/CUDA unavailable: {cuda_message}")

    for warning in warnings:
        print(f"WARNING: {warning}")

    if errors:
        print("1C training preflight failed:", file=sys.stderr)
        for error in errors:
            print(f"- {error}", file=sys.stderr)
        return 1

    print("1C training preflight passed.")
    return 0


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