import os
import json
import re
import pandas as pd
import comfy.samplers  # Taps into ComfyUI's native sampler/scheduler listings
import folder_paths  # Taps into ComfyUI's native model folder tracking

STYLES_DIR = os.path.join(os.path.dirname(__file__), "styles")
CONFIGS_DIR = os.path.join(os.path.dirname(__file__), "configs")
GEN_PRESETS_FILE = os.path.join(CONFIGS_DIR, "generation_presets.json")


def ensure_styles_dir():
    os.makedirs(STYLES_DIR, exist_ok=True)


def get_style_files() -> list[str]:
    ensure_styles_dir()
    files = sorted(
        f
        for f in os.listdir(STYLES_DIR)
        if f.lower().endswith((".csv", ".xlsx", ".xls"))
    )
    return files if files else ["none"]


def load_styles(filename: str) -> dict:
    """Return {name: {prompt, negative_prompt, description, category}}."""

    if not filename or filename == "none":
        return {}

    path = os.path.join(STYLES_DIR, filename)

    if not os.path.exists(path):
        return {}

    ext = os.path.splitext(filename)[1].lower()

    if ext == ".csv":
        df = pd.read_csv(path, encoding="utf-8-sig")
    elif ext in (".xlsx", ".xls"):
        df = pd.read_excel(path)
    else:
        return {}

    styles = {}

    for _, row in df.iterrows():
        name = str(row.get("name", "")).strip()
        if not name:
            continue

        styles[name] = {
            "prompt": str(row.get("prompt", "")).strip(),
            "negative_prompt": str(row.get("negative_prompt", "")).strip(),
            "description": str(row.get("description", "")).strip(),
            "category": str(row.get("category", "")).strip(),
        }

    return styles


def get_style_names(csv_filename: str) -> list[str]:
    names = list(load_styles(csv_filename).keys())
    return names if names else ["none"]


def apply_style_text(input_text: str, style_text: str) -> str:
    """
    If style_text contains {prompt}, replace it with input_text.
    Otherwise, append input_text to the end (comma-separated).
    """
    if not style_text:
        return input_text
    if re.search(r"(?i)\{prompt\}", style_text):
        return re.sub(r"(?i)\{prompt\}", input_text, style_text)
    if input_text:
        return f"{style_text}, {input_text}"
    return style_text


def ensure_configs_dir():
    os.makedirs(CONFIGS_DIR, exist_ok=True)


def get_generation_presets():
    ensure_configs_dir()
    if not os.path.exists(GEN_PRESETS_FILE):
        # Fetch available upscale models to set a valid default if possible
        models = folder_paths.get_filename_list("upscale_models")
        default_model = models[0] if models else "None"
        default_data = [
            {
                "name": "dreamshaper 8",
                "width": 512,
                "height": 768,
                "cfg": 8.0,
                "steps": 30,
                "sampler_name": "dpmpp_sde",
                "scheduler": "karras",
                "hires_steps": 15,
                "hires_cfg": 5.5,
                "hires_denoise": 0.4,
                "hires_upscale_by": 2.0,
                "hires_model": default_model,
                "is_latent": False,
            }
        ]
        with open(GEN_PRESETS_FILE, "w", encoding="utf-8") as f:
            json.dump(default_data, f, indent=2)
    presets_dict = {}
    try:
        with open(GEN_PRESETS_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                for item in data:
                    if "name" in item:
                        presets_dict[item["name"]] = item
    except Exception as e:
        print(f"[Generation Preset Error] Failed to read settings: {e}")
    return presets_dict