from __future__ import annotations

import argparse
import os
from pathlib import Path


DEFAULT_KEY_PATH = Path("KEY.md")


def load_key_file(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    loaded: dict[str, str] = {}
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        loaded[key.strip()] = value.strip()
    return loaded


def ensure_tokens(args: argparse.Namespace) -> None:
    primary_key_path = Path(getattr(args, "key_path", DEFAULT_KEY_PATH))
    if not primary_key_path.exists():
        parent_key = primary_key_path.parent.parent / primary_key_path.name
        if parent_key.exists():
            primary_key_path = parent_key
    values = load_key_file(primary_key_path)

    if not getattr(args, "hf_token", None):
        args.hf_token = os.getenv("HF_TOKEN") or values.get("HF_TOKEN")

    modelhub_token = getattr(args, "modelhub_token", None)
    if not modelhub_token:
        modelhub_token = os.getenv("MODELHUB_XC_TOKEN") or os.getenv("XC_TOKEN") or values.get("MODELHUB_XC_TOKEN") or values.get("XC_TOKEN")

    args.modelhub_token = modelhub_token

    if args.hf_token:
        os.environ["HF_TOKEN"] = args.hf_token
    if args.modelhub_token:
        os.environ["MODELHUB_XC_TOKEN"] = args.modelhub_token
        os.environ["XC_TOKEN"] = args.modelhub_token
    if not args.modelhub_token and not os.getenv("STRATEGY_ID"):
        raise ValueError("XC_TOKEN is required, either via environment or KEY.md")