from __future__ import annotations

import json
from pathlib import Path


TOOL_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = TOOL_DIR.parent

RUNS_DIR = TOOL_DIR / "runs"
OUTPUT_DIR = TOOL_DIR / "output"
SD_SCRIPTS_DIR = TOOL_DIR / "sd-scripts"
USER_SETTINGS_PATH = TOOL_DIR / "user_settings.json"
HISTORY_DIR = TOOL_DIR / "learning_history"
HISTORY_ASSETS_DIR = HISTORY_DIR / "assets"
HISTORY_PATH = HISTORY_DIR / "history.json"

TAGGER_DIRS = [
    TOOL_DIR / "models" / "tagger",
]

DEFAULT_TRAIN_SETTINGS = {
    "max_train_steps": 500,
    "learning_rate": 1e-4,
    "train_batch_size": 2,
    "network_dim": 16,
    "network_alpha": 16,
    "merge_strength": 1.0,
    "seed": 42,
}

DEFAULT_ANIMA_TRAIN_SETTINGS = {
    "max_train_steps": 600,
    "learning_rate": 1e-4,
    "train_batch_size": 1,
    "network_dim": 8,
    "network_alpha": 1,
    "merge_strength": 1.25,
    "seed": 42,
}

SPEED_PRESET_LABELS = [
    "安定学習セット",
    "通常学習セット",
    "詳細学習セット",
]

DEFAULT_SPEED_PRESET = "通常学習セット"

SPEED_PRESETS = {
    "SDXL": {
        "安定学習セット": {
            **DEFAULT_TRAIN_SETTINGS,
            "max_train_steps": 300,
            "learning_rate": 1e-4,
            "train_batch_size": 1,
            "network_dim": 16,
            "network_alpha": 16,
            "merge_strength": 1.0,
        },
        "通常学習セット": {
            **DEFAULT_TRAIN_SETTINGS,
            "max_train_steps": 500,
            "learning_rate": 1e-4,
            "train_batch_size": 2,
            "network_dim": 16,
            "network_alpha": 16,
            "merge_strength": 1.0,
        },
        "詳細学習セット": {
            **DEFAULT_TRAIN_SETTINGS,
            "max_train_steps": 600,
            "learning_rate": 8e-5,
            "train_batch_size": 2,
            "network_dim": 32,
            "network_alpha": 32,
            "merge_strength": 1.0,
        },
    },
    "Anima": {
        "安定学習セット": {
            **DEFAULT_ANIMA_TRAIN_SETTINGS,
            "max_train_steps": 200,
            "learning_rate": 8e-5,
            "train_batch_size": 1,
            "network_dim": 8,
            "network_alpha": 2,
            "merge_strength": 1.25,
        },
        "通常学習セット": {
            **DEFAULT_ANIMA_TRAIN_SETTINGS,
            "max_train_steps": 300,
            "learning_rate": 1.6e-4,
            "train_batch_size": 2,
            "network_dim": 8,
            "network_alpha": 1,
            "merge_strength": 1.25,
        },
        "詳細学習セット": {
            **DEFAULT_ANIMA_TRAIN_SETTINGS,
            "max_train_steps": 400,
            "learning_rate": 1.2e-4,
            "train_batch_size": 2,
            "network_dim": 16,
            "network_alpha": 8,
            "merge_strength": 1.25,
        },
    },
}

STYLE_TAGS = {
    "anime_style",
    "artstyle",
    "backlighting",
    "brush",
    "brush_stroke",
    "chibi",
    "colored_pencil",
    "colorful",
    "cross_hatching",
    "dramatic_lighting",
    "flat_color",
    "flat_shading",
    "gradient",
    "greyscale",
    "halftone",
    "high_contrast",
    "lineart",
    "marker",
    "monochrome",
    "oil_painting",
    "painting",
    "pastel",
    "painterly",
    "photorealistic",
    "realistic",
    "semi-realistic",
    "rough_sketch",
    "sketch",
    "soft_lighting",
    "screenshot",
    "anime_screenshot",
    "super_deformed",
    "traditional_media",
    "watercolor",
}

NOISY_CAPTION_TAGS = {
    "absurdres",
    "artist_name",
    "bad_anatomy",
    "bad_hands",
    "bad_id",
    "bad_pixiv_id",
    "blurry",
    "censored",
    "commentary",
    "compression_artifacts",
    "error",
    "extra_digits",
    "fewer_digits",
    "highres",
    "jpeg_artifacts",
    "logo",
    "lowres",
    "missing_fingers",
    "multiple_views",
    "normal_quality",
    "out_of_frame",
    "patreon_username",
    "pixiv_id",
    "poorly_drawn",
    "scan_artifacts",
    "signature",
    "speech_bubble",
    "text",
    "translation_request",
    "username",
    "watermark",
    "web_address",
    "worst_quality",
}

NOISY_TAG_PREFIXES = (
    "artist:",
    "copyright:",
    "character:",
    "commentary_request",
    "commissioner:",
    "md5_mismatch",
    "source_",
)

NOISY_TAG_PARTS = (
    "watermark",
    "signature",
    "username",
    "pixiv",
    "twitter",
    "patreon",
    "deviantart",
    "sample_watermark",
)


def default_user_settings() -> dict[str, str]:
    models_root = TOOL_DIR / "models"
    common_model_dir = str(models_root / "StableDiffusion")
    return {
        "models_root_dir": str(models_root),
        "lora_output_dir": str(OUTPUT_DIR),
        "sdxl_model_dir": common_model_dir,
        "anima_model_dir": common_model_dir,
        "anima_text_encoder_dir": str(models_root / "TextEncoders"),
        "anima_vae_dir": str(models_root / "VAE"),
        "anima_adapter_dir": str(models_root / "Anima"),
        "anima_dit_path": "",
        "anima_qwen3_path": "",
        "anima_vae_path": "",
        "anima_llm_adapter_path": "",
        "last_lora_output_name": "mylora",
    }


def load_user_settings() -> dict[str, str]:
    settings = default_user_settings()
    if USER_SETTINGS_PATH.exists():
        try:
            loaded = json.loads(USER_SETTINGS_PATH.read_text(encoding="utf-8"))
            if isinstance(loaded, dict):
                settings.update({str(k): str(v) for k, v in loaded.items()})
        except Exception:
            pass
    return settings


def save_user_settings(values: dict[str, str]) -> dict[str, str]:
    settings = load_user_settings()
    settings.update({str(k): str(v) for k, v in values.items()})
    USER_SETTINGS_PATH.write_text(
        json.dumps(settings, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    return settings


def normalize_user_path(value: str | Path | None, fallback: Path) -> Path:
    text = str(value or fallback).strip().strip('"').strip("'")
    if not text:
        text = str(fallback)
    return Path(text).expanduser()


def get_lora_output_dir() -> Path:
    value = load_user_settings().get("lora_output_dir") or str(OUTPUT_DIR)
    return normalize_user_path(value, OUTPUT_DIR)


def infer_models_root_dir(user_settings: dict[str, str] | None = None) -> Path:
    values = user_settings or load_user_settings()
    explicit = values.get("models_root_dir", "")
    explicit_path = normalize_user_path(explicit, TOOL_DIR / "models")
    default_root = TOOL_DIR / "models"
    if explicit and explicit_path.exists() and explicit_path != default_root:
        return explicit_path

    for key in ("sdxl_model_dir", "anima_model_dir", "anima_text_encoder_dir", "anima_vae_dir"):
        raw = values.get(key, "")
        if not raw:
            continue
        path = normalize_user_path(raw, TOOL_DIR / "models")
        name = path.name.lower()
        if name in {
            "stablediffusion",
            "stable-diffusion",
            "stable_diffusion",
            "checkpoints",
            "sdxl",
            "textencoders",
            "text_encoders",
            "text-encoders",
            "vae",
            "vaeapprox",
        }:
            return path.parent
        if path.exists() and any((path / child).exists() for child in ("StableDiffusion", "TextEncoders", "VAE", "Vae")):
            return path
    return explicit_path


def get_models_root_dir() -> Path:
    return infer_models_root_dir(load_user_settings())


def get_sdxl_model_dir() -> Path:
    value = load_user_settings().get("sdxl_model_dir") or str(TOOL_DIR / "models" / "StableDiffusion")
    return normalize_user_path(value, TOOL_DIR / "models" / "StableDiffusion")


def get_anima_model_dir() -> Path:
    user_settings = load_user_settings()
    value = user_settings.get("anima_model_dir") or user_settings.get("sdxl_model_dir") or str(TOOL_DIR / "models" / "StableDiffusion")
    return normalize_user_path(value, TOOL_DIR / "models" / "StableDiffusion")


def get_anima_text_encoder_dir() -> Path:
    value = load_user_settings().get("anima_text_encoder_dir") or str(TOOL_DIR / "models" / "TextEncoders")
    return normalize_user_path(value, TOOL_DIR / "models" / "TextEncoders")


def get_anima_vae_dir() -> Path:
    value = load_user_settings().get("anima_vae_dir") or str(TOOL_DIR / "models" / "VAE")
    return normalize_user_path(value, TOOL_DIR / "models" / "VAE")


def get_anima_adapter_dir() -> Path:
    fallback = TOOL_DIR / "models" / "Anima"
    value = load_user_settings().get("anima_adapter_dir") or str(fallback)
    path = normalize_user_path(value, fallback)
    return path if path.exists() else fallback
