"""Load IsolatedEnv configuration from TOML files.

This module provides declarative configuration for isolated environments,
allowing custom nodes to define their requirements in a TOML file instead
of programmatically.

Config file: comfy-env.toml

Simplified config (recommended for CUDA packages only):

    # comfy-env.toml - just list CUDA packages
    [cuda]
    torch-scatter = "2.1.2"
    torch-cluster = "1.6.3"
    spconv = "*"  # latest compatible

Or as a list:

    cuda = [
        "torch-scatter==2.1.2",
        "torch-cluster==1.6.3",
        "spconv",
    ]

Optional overrides:

    [env]
    cuda = "12.4"       # Override auto-detection
    pytorch = "2.5.1"   # Override auto-detection

Full format:

    [env]
    name = "my-node"
    python = "3.10"
    cuda = "auto"

    [packages]
    requirements = ["my-package>=1.0.0"]
    no_deps = ["torch-scatter==2.1.2"]

    [sources]
    wheel_sources = ["https://my-wheels.github.io/"]

Available auto-derived variables:
    - {cuda_version}: Full CUDA version (e.g., "12.8")
    - {cuda_short}: CUDA version without dot (e.g., "128")
    - {pytorch_version}: Full PyTorch version (e.g., "2.9.1")
    - {pytorch_short}: PyTorch version without dots (e.g., "291")
    - {pytorch_mm}: PyTorch major.minor without dot (e.g., "29")
"""

import sys
from pathlib import Path
from typing import Optional, Dict, Any, List

# Use built-in tomllib (Python 3.11+) or tomli fallback
if sys.version_info >= (3, 11):
    import tomllib
else:
    try:
        import tomli as tomllib
    except ImportError:
        tomllib = None  # type: ignore

from .config import IsolatedEnv, EnvManagerConfig, LocalConfig, NodeReq, ToolConfig
from .cuda_gpu_detection import detect_cuda_version


# Config file name
CONFIG_FILE_NAMES = [
    "comfy-env.toml",
]


def load_env_from_file(
    path: Path,
    base_dir: Optional[Path] = None,
) -> IsolatedEnv:
    """
    Load IsolatedEnv configuration from a TOML file.

    Args:
        path: Path to the TOML config file
        base_dir: Base directory for resolving relative paths (default: file's parent)

    Returns:
        Configured IsolatedEnv instance

    Raises:
        FileNotFoundError: If config file doesn't exist
        ValueError: If config is invalid
        ImportError: If tomli is not installed (Python < 3.11)

    Example:
        >>> env = load_env_from_file(Path("my_node/comfy-env.toml"))
        >>> print(env.name)
        'my-node'
    """
    if tomllib is None:
        raise ImportError(
            "TOML parsing requires tomli for Python < 3.11. "
            "Install it with: pip install tomli"
        )

    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {path}")

    base_dir = Path(base_dir) if base_dir else path.parent

    with open(path, "rb") as f:
        data = tomllib.load(f)

    return _parse_config(data, base_dir)


def discover_env_config(
    node_dir: Path,
    file_names: Optional[List[str]] = None,
) -> Optional[IsolatedEnv]:
    """
    Auto-discover and load config from a node directory.

    Searches for standard config file names in order of priority.

    Args:
        node_dir: Directory to search for config files
        file_names: Custom list of file names to search (default: CONFIG_FILE_NAMES)

    Returns:
        IsolatedEnv if config found, None otherwise

    Example:
        >>> env = discover_env_config(Path("my_custom_node/"))
        >>> if env:
        ...     print(f"Found config: {env.name}")
        ... else:
        ...     print("No config file found")
    """
    if tomllib is None:
        # Can't parse TOML without the library
        return None

    node_dir = Path(node_dir)
    file_names = file_names or CONFIG_FILE_NAMES

    for name in file_names:
        config_path = node_dir / name
        if config_path.exists():
            return load_env_from_file(config_path, node_dir)

    return None


def _get_default_pytorch_version(cuda_version: Optional[str]) -> str:
    """
    Get default PyTorch version based on CUDA version.

    Args:
        cuda_version: CUDA version (e.g., "12.4", "12.8") or None

    Returns:
        PyTorch version string

    Version Mapping:
        - CUDA 12.4 (Pascal): PyTorch 2.5.1
        - CUDA 12.8 (Turing+): PyTorch 2.8.0
    """
    if cuda_version == "12.4":
        return "2.5.1"  # Full: Pascal GPUs
    return "2.8.0"  # Modern: Turing through Blackwell


def _parse_config(data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
    """
    Parse TOML data into IsolatedEnv.

    Supports both simplified and full config formats:

    Simplified (CUDA packages only):
        [packages]
        torch-scatter = "2.1.2"
        torch-cluster = "1.6.3"

    Or as list:
        packages = ["torch-scatter==2.1.2", "torch-cluster==1.6.3"]

    Full:
        [env]
        name = "my-node"
        [packages]
        no_deps = ["torch-scatter==2.1.2"]

    Args:
        data: Parsed TOML data
        base_dir: Base directory for resolving relative paths

    Returns:
        Configured IsolatedEnv instance
    """
    env_section = data.get("env", {})
    packages_section = data.get("packages", {})
    sources_section = data.get("sources", {})
    worker_section = data.get("worker", {})
    variables = dict(data.get("variables", {}))  # Copy to avoid mutation

    # Handle CUDA version - default to "auto" if not specified
    cuda = env_section.get("cuda", "auto")
    if cuda == "auto":
        cuda = detect_cuda_version()
    elif cuda == "null" or cuda == "none":
        cuda = None

    # Add auto-derived variables based on CUDA
    if cuda:
        variables.setdefault("cuda_version", cuda)
        variables.setdefault("cuda_short", cuda.replace(".", ""))

    # Handle pytorch version - auto-derive if "auto" or not specified
    pytorch_version = env_section.get("pytorch_version") or env_section.get("pytorch")
    if pytorch_version == "auto" or (pytorch_version is None and cuda):
        pytorch_version = _get_default_pytorch_version(cuda)

    if pytorch_version:
        variables.setdefault("pytorch_version", pytorch_version)
        # Add short version without dots (e.g., "2.9.1" -> "291")
        pytorch_short = pytorch_version.replace(".", "")
        variables.setdefault("pytorch_short", pytorch_short)
        # Add major.minor without dot (e.g., "2.9.1" -> "29") for wheel naming
        parts = pytorch_version.split(".")[:2]
        pytorch_mm = "".join(parts)
        variables.setdefault("pytorch_mm", pytorch_mm)

    # Parse CUDA packages - support multiple formats
    # Priority: [cuda] section > cuda = [...] > [packages] section
    no_deps_requirements = []
    requirements = []

    cuda_section = data.get("cuda", {})

    if cuda_section:
        # New format: [cuda] section or cuda = [...]
        if isinstance(cuda_section, list):
            # Format: cuda = ["torch-scatter==2.1.2", ...]
            no_deps_requirements = [_substitute_vars(req, variables) for req in cuda_section]
        elif isinstance(cuda_section, dict):
            # Format: [cuda] with package = "version" pairs
            for pkg, ver in cuda_section.items():
                if ver == "*" or ver == "":
                    no_deps_requirements.append(pkg)
                else:
                    no_deps_requirements.append(f"{pkg}=={ver}")

    elif isinstance(packages_section, list):
        # Legacy format: packages = ["torch-scatter==2.1.2", ...]
        no_deps_requirements = [_substitute_vars(req, variables) for req in packages_section]

    elif isinstance(packages_section, dict):
        # Check for simplified format: [packages] with key=value pairs
        # vs old format: [packages] with requirements/no_deps lists

        has_old_keys = any(k in packages_section for k in ["requirements", "no_deps", "requirements_file"])

        if has_old_keys:
            # Old format
            raw_requirements = packages_section.get("requirements", [])
            requirements = [_substitute_vars(req, variables) for req in raw_requirements]

            raw_no_deps = packages_section.get("no_deps", [])
            no_deps_requirements = [_substitute_vars(req, variables) for req in raw_no_deps]
        else:
            # Simplified format: [packages] with package = "version" pairs
            # All packages are CUDA packages
            for pkg, ver in packages_section.items():
                if ver == "*" or ver == "":
                    no_deps_requirements.append(pkg)
                else:
                    no_deps_requirements.append(f"{pkg}=={ver}")

    # Resolve requirements_file path (relative to base_dir)
    requirements_file = None
    if isinstance(packages_section, dict) and "requirements_file" in packages_section:
        req_file_path = packages_section["requirements_file"]
        requirements_file = base_dir / req_file_path

    # Get wheel sources and index URLs (optional - registry handles most cases now)
    wheel_sources = sources_section.get("wheel_sources", [])
    index_urls = sources_section.get("index_urls", [])

    # Parse worker configuration
    worker_package = worker_section.get("package")
    worker_script = worker_section.get("script")

    return IsolatedEnv(
        name=env_section.get("name", base_dir.name),
        python=env_section.get("python", "3.10"),
        cuda=cuda,
        pytorch_version=pytorch_version,
        requirements=requirements,
        no_deps_requirements=no_deps_requirements,
        requirements_file=requirements_file,
        wheel_sources=wheel_sources,
        index_urls=index_urls,
        worker_package=worker_package,
        worker_script=worker_script,
    )


def _substitute_vars(s: str, variables: Dict[str, str]) -> str:
    """
    Substitute {var_name} placeholders with values from variables dict.

    Args:
        s: String with placeholders like {var_name}
        variables: Dictionary mapping variable names to values

    Returns:
        String with placeholders replaced

    Example:
        >>> _substitute_vars("torch=={pytorch_version}", {"pytorch_version": "2.4.1"})
        'torch==2.4.1'
    """
    for key, value in variables.items():
        s = s.replace(f"{{{key}}}", str(value))
    return s


# =============================================================================
# V2 Schema Parser
# =============================================================================

# Reserved table names that are NOT isolated environments
RESERVED_TABLES = {"local", "node_reqs", "env", "packages", "sources", "cuda", "variables", "worker", "tools"}


def load_config(
    path: Path,
    base_dir: Optional[Path] = None,
) -> EnvManagerConfig:
    """
    Load full EnvManagerConfig from a TOML file.

    Supports both full schema (named envs) and simple format (auto-detected).

    Args:
        path: Path to the TOML config file
        base_dir: Base directory for resolving relative paths

    Returns:
        EnvManagerConfig with local, envs, and node_reqs

    Example:
        >>> config = load_config(Path("comfy-env.toml"))
        >>> if config.has_local:
        ...     print(f"Local CUDA packages: {config.local.cuda_packages}")
        >>> for name, env in config.envs.items():
        ...     print(f"Isolated env: {name}")
    """
    if tomllib is None:
        raise ImportError(
            "TOML parsing requires tomli for Python < 3.11. "
            "Install it with: pip install tomli"
        )

    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Config file not found: {path}")

    base_dir = Path(base_dir) if base_dir else path.parent

    with open(path, "rb") as f:
        data = tomllib.load(f)

    return _parse_full_config(data, base_dir)


def discover_config(
    node_dir: Path,
    file_names: Optional[List[str]] = None,
) -> Optional[EnvManagerConfig]:
    """
    Auto-discover and load EnvManagerConfig from a node directory.

    Args:
        node_dir: Directory to search for config files
        file_names: Custom list of file names to search

    Returns:
        EnvManagerConfig if found, None otherwise
    """
    if tomllib is None:
        return None

    node_dir = Path(node_dir)
    file_names = file_names or CONFIG_FILE_NAMES

    for name in file_names:
        config_path = node_dir / name
        if config_path.exists():
            return load_config(config_path, node_dir)

    return None


def _parse_full_config(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
    """
    Parse TOML data into EnvManagerConfig.

    Schema:
        [local.cuda]        - CUDA packages for host
        [local.packages]    - Regular packages for host
        [envname]           - Isolated env (python, cuda, pytorch)
        [envname.cuda]      - CUDA packages for env
        [envname.packages]  - Packages for env
        [node_reqs]         - Node dependencies

    Also supports simple format ([env] + [packages]) for backward compatibility.
    """
    # Detect if this is full schema or simple format
    is_full = "local" in data or _has_named_env(data)

    if not is_full:
        # Simple format - convert to full structure
        return _convert_simple_to_full(data, base_dir)

    # Parse full schema
    local = _parse_local_section(data.get("local", {}))
    envs = _parse_env_sections(data, base_dir)
    node_reqs = _parse_node_reqs(data.get("node_reqs", {}))
    tools = _parse_tools_section(data.get("tools", {}))

    return EnvManagerConfig(
        local=local,
        envs=envs,
        node_reqs=node_reqs,
        tools=tools,
    )


def _has_named_env(data: Dict[str, Any]) -> bool:
    """Check if data has any named environment tables (not reserved names)."""
    env_indicators = ["python", "cuda", "pytorch", "cuda_version", "pytorch_version"]
    for key in data.keys():
        if key not in RESERVED_TABLES and isinstance(data[key], dict):
            # Check if it looks like an env definition
            section = data[key]
            if any(k in section for k in env_indicators):
                return True
    return False


def _parse_local_section(local_data: Dict[str, Any]) -> LocalConfig:
    """Parse [local] section."""
    cuda_packages = {}
    requirements = []

    # [local.cuda] - CUDA packages
    cuda_section = local_data.get("cuda", {})
    if isinstance(cuda_section, dict):
        for pkg, ver in cuda_section.items():
            cuda_packages[pkg] = ver if ver and ver != "*" else ""

    # [local.packages] - regular packages
    packages_section = local_data.get("packages", {})
    if isinstance(packages_section, dict):
        requirements = packages_section.get("requirements", [])
    elif isinstance(packages_section, list):
        requirements = packages_section

    return LocalConfig(
        cuda_packages=cuda_packages,
        requirements=requirements,
    )


def _parse_env_sections(data: Dict[str, Any], base_dir: Path) -> Dict[str, IsolatedEnv]:
    """Parse named environment sections."""
    envs = {}
    env_indicators = ["python", "cuda", "pytorch", "cuda_version", "pytorch_version"]

    for key, value in data.items():
        if key in RESERVED_TABLES:
            continue
        if not isinstance(value, dict):
            continue

        # Check if this looks like an env definition
        if not any(k in value for k in env_indicators):
            continue

        env = _parse_single_env(key, value, base_dir)
        envs[key] = env

    return envs


def _parse_single_env(name: str, env_data: Dict[str, Any], base_dir: Path) -> IsolatedEnv:
    """Parse a single isolated environment section."""
    # Get basic env config
    python = env_data.get("python", "3.10")
    # Support both "cuda" and "cuda_version" field names (cuda_version avoids conflict with [envname.cuda] table)
    cuda = env_data.get("cuda_version") or env_data.get("cuda", "auto")
    pytorch = env_data.get("pytorch_version") or env_data.get("pytorch", "auto")

    # Handle auto-detection
    if cuda == "auto":
        cuda = detect_cuda_version()
    elif cuda in ("null", "none", None):
        cuda = None

    if pytorch == "auto":
        pytorch = _get_default_pytorch_version(cuda)

    # Parse [envname.cuda] - CUDA packages
    cuda_section = env_data.get("cuda", {})
    no_deps_requirements = []
    if isinstance(cuda_section, dict):
        for pkg, ver in cuda_section.items():
            if ver == "*" or ver == "":
                no_deps_requirements.append(pkg)
            else:
                no_deps_requirements.append(f"{pkg}=={ver}")

    # Parse [envname.packages] - regular packages
    packages_section = env_data.get("packages", {})
    requirements = []
    if isinstance(packages_section, dict):
        requirements = packages_section.get("requirements", [])
    elif isinstance(packages_section, list):
        requirements = packages_section

    # Parse platform-specific packages [envname.packages.windows], etc.
    windows_reqs = []
    linux_reqs = []
    darwin_reqs = []

    if isinstance(packages_section, dict):
        win_section = packages_section.get("windows", {})
        if isinstance(win_section, dict):
            windows_reqs = win_section.get("requirements", [])
        elif isinstance(win_section, list):
            windows_reqs = win_section

        linux_section = packages_section.get("linux", {})
        if isinstance(linux_section, dict):
            linux_reqs = linux_section.get("requirements", [])
        elif isinstance(linux_section, list):
            linux_reqs = linux_section

        darwin_section = packages_section.get("darwin", {})
        if isinstance(darwin_section, dict):
            darwin_reqs = darwin_section.get("requirements", [])
        elif isinstance(darwin_section, list):
            darwin_reqs = darwin_section

    return IsolatedEnv(
        name=name,
        python=python,
        cuda=cuda,
        pytorch_version=pytorch,
        requirements=requirements,
        no_deps_requirements=no_deps_requirements,
        windows_requirements=windows_reqs,
        linux_requirements=linux_reqs,
        darwin_requirements=darwin_reqs,
    )


def _parse_node_reqs(node_reqs_data: Dict[str, Any]) -> List[NodeReq]:
    """Parse [node_reqs] section."""
    reqs = []

    for name, value in node_reqs_data.items():
        if isinstance(value, str):
            # Simple format: VideoHelperSuite = "Kosinkadink/ComfyUI-VideoHelperSuite"
            reqs.append(NodeReq(name=name, repo=value))
        elif isinstance(value, dict):
            # Extended format: VideoHelperSuite = { repo = "..." }
            repo = value.get("repo", "")
            reqs.append(NodeReq(name=name, repo=repo))

    return reqs


def _parse_tools_section(tools_data: Dict[str, Any]) -> Dict[str, ToolConfig]:
    """Parse [tools] section.

    Supports:
        [tools]
        blender = "4.2"

    Or extended:
        [tools.blender]
        version = "4.2"
        install_dir = "/custom/path"
    """
    tools = {}

    for name, value in tools_data.items():
        if isinstance(value, str):
            # Simple format: blender = "4.2"
            tools[name] = ToolConfig(name=name, version=value)
        elif isinstance(value, dict):
            # Extended format: [tools.blender] with version, install_dir
            version = value.get("version", "latest")
            install_dir = value.get("install_dir")
            if install_dir:
                install_dir = Path(install_dir)
            tools[name] = ToolConfig(name=name, version=version, install_dir=install_dir)

    return tools


def _convert_simple_to_full(data: Dict[str, Any], base_dir: Path) -> EnvManagerConfig:
    """Convert simple config format to full EnvManagerConfig.

    Simple configs have [env] and [packages] sections but no named environments.
    This converts them to the full format with a single named environment.
    """
    # Parse using simple parser to get IsolatedEnv
    simple_env = _parse_config(data, base_dir)

    # Parse tools section (shared between simple and full format)
    tools = _parse_tools_section(data.get("tools", {}))

    # Check if this has explicit env settings (isolated venv) vs just CUDA packages (local install)
    env_section = data.get("env", {})
    has_explicit_env = bool(env_section.get("name") or env_section.get("python"))

    if has_explicit_env:
        # Isolated venv config
        return EnvManagerConfig(
            local=LocalConfig(),
            envs={simple_env.name: simple_env},
            node_reqs=[],
            tools=tools,
        )
    else:
        # Local CUDA packages only (no isolated venv)
        cuda_packages = {}
        for req in simple_env.no_deps_requirements:
            if "==" in req:
                pkg, ver = req.split("==", 1)
                cuda_packages[pkg] = ver
            else:
                cuda_packages[req] = ""

        return EnvManagerConfig(
            local=LocalConfig(
                cuda_packages=cuda_packages,
                requirements=simple_env.requirements,
            ),
            envs={},
            node_reqs=[],
            tools=tools,
        )