from pathlib import Path
import toml # type: ignore
from typing import Any


APP_NAME = "comfy-resolve"
def get_config_path() -> Path:
    """Returns the path to the configuration file."""
    return Path.home() / ".config" / APP_NAME / "config.toml"

def config_exists() -> bool:
    """Checks if the configuration file exists."""
    return get_config_path().is_file()

def load_config() -> dict[str, Any]:
    """Loads the configuration from the TOML file."""
    if not config_exists():
        return {}
    with open(get_config_path(), "r") as f:
        return toml.load(f) # type: ignore

def infer_model_paths(comfyui_dir: str) -> dict[str, Path]:
    """Infers model paths based on the standard ComfyUI folder layout."""
    base_path = Path(comfyui_dir) / "models"
    return {
        "checkpoints": base_path / "checkpoints",
        "loras": base_path / "loras",
        "vae": base_path / "vae",
        "text_encoders": base_path / "text_encoders",
        "controlnet": base_path / "controlnet",
        "upscale_models": base_path / "upscale_models",
        "embeddings": base_path / "embeddings",
        "clip": base_path / "clip",
        "ipadapter": base_path / "ipadapter",
        "gligen": base_path / "gligen",
        "style_models": base_path / "style_models",
        "hypernetworks": base_path / "hypernetworks",
        "diffusers": base_path / "diffusers",
        "unclip": base_path / "unclip",
        "insightface": base_path / "insightface",
        "photomaker": base_path / "photomaker",
        "sadtalker": base_path / "sadtalker",
        "audio": base_path / "audio",
        "other": base_path / "other",
    }

def write_default_config(comfyui_dir: str) -> Path:
    """Writes the default config template with inferred model paths."""
    config_path = get_config_path()
    config_path.parent.mkdir(parents=True, exist_ok=True)

    inferred_paths = infer_model_paths(comfyui_dir)
    paths_section = "\n".join([f"# {k} = \"{v.as_posix()}\"" for k, v in inferred_paths.items()])

    default_content = f"""
comfyui_dir = "{Path(comfyui_dir).as_posix()}"

# Individual paths are inferred from comfyui_dir by default.
# Uncomment and edit any line below to override a specific path.

# [paths]
{paths_section}

[api_keys]
# Can also be set via environment variables: COMFY_HF_TOKEN, COMFY_CIVITAI_TOKEN
huggingface = ""
civitai = ""

[download]
concurrent = 2
max_retries = 3

[resolution]
preferred_source = "huggingface"   # "huggingface" | "civitai" | "any"
auto_confidence_threshold = "high" # "high" | "medium"

[mappings]
# Static overrides for private or custom models
# "my_custom_lora.safetensors" = {{ type = "direct_url", url = "https://..." }}
"""
    with open(config_path, "w") as f:
        f.write(default_content.strip())
    return config_path

def get_model_path(asset_type: str, config: dict[str, Any]) -> Path:
    """Returns the override path if set in [paths], else inferred path from comfyui_dir."""
    comfyui_dir = config.get("comfyui_dir")
    if not comfyui_dir:
        raise ValueError("comfyui_dir not found in config.")

    if "paths" in config and asset_type in config["paths"]:
        return Path(config["paths"][asset_type])
    
    inferred_paths = infer_model_paths(comfyui_dir)
    return inferred_paths.get(asset_type, Path(comfyui_dir) / "models" / "other")
