from __future__ import annotations

import csv
import math
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

from PIL import Image, ImageOps, ImageStat

import settings

DEFAULT_GENERAL_TAG_THRESHOLD = 0.35
DETAILED_GENERAL_TAG_THRESHOLD = 0.22
DEFAULT_CHARACTER_TAG_THRESHOLD = 0.5
DEFAULT_MAX_CONTENT_TAGS = 36
DETAILED_MAX_CONTENT_TAGS = 72


@dataclass
class CaptionAnalysis:
    tags: list[str]
    caption: str
    method: str
    notes: list[str]
    traits: dict[str, str | float]


def clean_tag(tag: str) -> str:
    return tag.strip().replace(" ", "_").lower()


def dedupe(tags: Iterable[str]) -> list[str]:
    seen: set[str] = set()
    result: list[str] = []
    for tag in tags:
        tag = clean_tag(tag)
        if not tag or tag in seen:
            continue
        seen.add(tag)
        result.append(tag)
    return result


def is_allowed_caption_tag(tag: str) -> bool:
    tag = clean_tag(tag)
    if not tag:
        return False
    if tag in settings.STYLE_TAGS or tag in settings.NOISY_CAPTION_TAGS:
        return False
    if any(tag.startswith(prefix) for prefix in settings.NOISY_TAG_PREFIXES):
        return False
    if any(part in tag for part in settings.NOISY_TAG_PARTS):
        return False
    return True


def sanitize_tags(tags: Iterable[str]) -> list[str]:
    return [tag for tag in dedupe(tags) if is_allowed_caption_tag(tag)]


def image_traits(image_path: str | Path) -> dict[str, str | float]:
    image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
    small = image.resize((128, 128), Image.Resampling.BICUBIC)
    hsv = small.convert("HSV")
    stat = ImageStat.Stat(hsv)
    saturation = float(stat.mean[1])
    value = float(stat.mean[2])

    gray = small.convert("L")
    gray_stat = ImageStat.Stat(gray)
    contrast = float(gray_stat.stddev[0])

    monochrome = saturation < 28
    bright_background = value > 180
    line_like = monochrome and bright_background and contrast > 45
    aspect = image.width / max(image.height, 1)

    return {
        "width": image.width,
        "height": image.height,
        "aspect": round(aspect, 3),
        "saturation": round(saturation, 2),
        "brightness": round(value, 2),
        "contrast": round(contrast, 2),
        "color_mode": "monochrome" if monochrome else "color",
        "line_mode": "lineart" if line_like else "painted",
    }


class WdTagger:
    def __init__(self, model_dir: Path):
        self.model_dir = model_dir
        self.model_path = model_dir / "model.onnx"
        self.tags_path = model_dir / "selected_tags.csv"
        self.config_path = model_dir / "config.json"
        self._session = None
        self._tags: list[tuple[str, int]] | None = None

    @property
    def available(self) -> bool:
        return self.model_path.exists() and self.tags_path.exists()

    def _load(self):
        try:
            import numpy as np
            import onnxruntime as ort
        except Exception as exc:  # pragma: no cover - optional runtime
            raise RuntimeError(f"WD tagger requires numpy and onnxruntime: {exc}") from exc

        if self._session is None:
            # A single 448px caption pass is fast enough on CPU. Using CPU also
            # avoids noisy CUDA DLL errors on machines that can train with
            # PyTorch but do not have ONNX Runtime CUDA dependencies installed.
            providers = ["CPUExecutionProvider"]
            self._session = ort.InferenceSession(str(self.model_path), providers=providers)

        if self._tags is None:
            tags: list[tuple[str, int]] = []
            with self.tags_path.open("r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    tags.append((clean_tag(row.get("name", "")), int(row.get("category", "0"))))
            self._tags = tags

        return np

    def analyze(
        self,
        image_path: str | Path,
        general_threshold: float = DEFAULT_GENERAL_TAG_THRESHOLD,
        character_threshold: float = DEFAULT_CHARACTER_TAG_THRESHOLD,
    ) -> list[str]:
        np = self._load()
        image = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
        size = 448
        canvas = Image.new("RGB", (max(image.size), max(image.size)), (255, 255, 255))
        canvas.paste(image, ((canvas.width - image.width) // 2, (canvas.height - image.height) // 2))
        canvas = canvas.resize((size, size), Image.Resampling.BICUBIC)

        array = np.asarray(canvas, dtype=np.float32)
        array = array[:, :, ::-1]
        array = np.expand_dims(array, 0)

        assert self._session is not None
        input_name = self._session.get_inputs()[0].name
        probs = self._session.run(None, {input_name: array})[0][0]

        selected: list[str] = []
        ignored_categories = {1, 3, 4, 9}
        for (tag, category), prob in zip(self._tags or [], probs):
            threshold = character_threshold if category == 4 else general_threshold
            if category in ignored_categories:
                continue
            if float(prob) >= threshold and is_allowed_caption_tag(tag):
                selected.append(tag)
        return selected


def find_tagger() -> WdTagger | None:
    for model_dir in settings.TAGGER_DIRS:
        tagger = WdTagger(model_dir)
        if tagger.available:
            return tagger
    return None


def heuristic_tags(traits: dict[str, str | float]) -> list[str]:
    tags = ["solo", "simple_background"]
    if traits["color_mode"] == "monochrome":
        tags.extend(["monochrome", "greyscale"])
    else:
        tags.extend(["color", "brown_hair", "brown_eyes"])

    if traits["line_mode"] == "lineart":
        tags.append("lineart")

    aspect = float(traits["aspect"])
    if 0.7 <= aspect <= 1.35:
        tags.extend(["portrait", "upper_body"])
    elif aspect < 0.7:
        tags.append("full_body")
    else:
        tags.append("wide_shot")
    return tags


def content_tags(tags: Iterable[str], max_tags: int = DEFAULT_MAX_CONTENT_TAGS) -> list[str]:
    return sanitize_tags(tags)[:max(1, int(max_tags or DEFAULT_MAX_CONTENT_TAGS))]


def make_caption(tags: Iterable[str], max_tags: int = DEFAULT_MAX_CONTENT_TAGS) -> str:
    tags = content_tags(tags, max_tags=max_tags)
    if not tags:
        tags = ["solo", "simple_background", "portrait"]

    return ", ".join(tags)


def clean_caption_reference(caption: str | None) -> str:
    if not caption:
        return ""
    return make_caption(part for part in caption.split(","))


def format_caption_for_display(caption: str | None) -> str:
    if not caption:
        return ""
    formatted: list[str] = []
    for tag in caption.split(","):
        tag = tag.strip().replace("_", " ")
        tag = re.sub(r"\s+", " ", tag).strip()
        if tag:
            formatted.append(tag)
    return ", ".join(formatted)


def format_caption_for_training(caption: str | None, model_type: str | None = "SDXL") -> str:
    if not caption:
        return ""
    tags = [tag.strip() for tag in caption.split(",") if tag.strip()]
    if (model_type or "").upper() != "ANIMA":
        return ", ".join(tags)

    return format_caption_for_display(caption)


def analyze_image(
    image_path: str | Path,
    general_threshold: float = DEFAULT_GENERAL_TAG_THRESHOLD,
    character_threshold: float = DEFAULT_CHARACTER_TAG_THRESHOLD,
    max_tags: int = DEFAULT_MAX_CONTENT_TAGS,
) -> CaptionAnalysis:
    traits = image_traits(image_path)
    notes: list[str] = []
    tagger = find_tagger()

    if tagger:
        try:
            tags = sanitize_tags(
                tagger.analyze(
                    image_path,
                    general_threshold=general_threshold,
                    character_threshold=character_threshold,
                )
            )
            method = "WD tagger"
            if not tags:
                notes.append("タグ推定モデルは見つかりましたが、しきい値を超える内容タグがありませんでした。")
                tags = sanitize_tags(heuristic_tags(traits))
                method = "画像特徴から推定"
        except Exception as exc:
            notes.append(f"WD taggerを使えなかったため、画像特徴から推定しました: {exc}")
            tags = sanitize_tags(heuristic_tags(traits))
            method = "画像特徴から推定"
    else:
        notes.append("models/tagger が見つからないため、簡易推定でキャプションを作りました。install.ps1 の -DownloadTagger を使うと精度が上がります。")
        tags = sanitize_tags(heuristic_tags(traits))
        method = "画像特徴から推定"

    caption = make_caption(tags, max_tags=max_tags)
    return CaptionAnalysis(tags=dedupe(tags), caption=caption, method=method, notes=notes, traits=traits)


def pair_similarity_notes(style: CaptionAnalysis, pair: CaptionAnalysis | None) -> list[str]:
    notes: list[str] = []
    if pair is None:
        notes.append("比較画像をまだ読み込んでいません。画像生成AIで比較画像を作り、下の比較画像欄へ入れてください。")
        return notes

    style_content = set(content_tags(style.tags))
    pair_content = set(content_tags(pair.tags))
    if style_content and pair_content:
        overlap = len(style_content & pair_content) / max(1, math.sqrt(len(style_content) * len(pair_content)))
        if overlap < 0.45:
            notes.append("2枚の自動タグがかなり違います。比較画像は構図・人物・服・背景をできるだけ同じにしてください。")

    aspect_delta = abs(float(style.traits["aspect"]) - float(pair.traits["aspect"]))
    if aspect_delta > 0.02:
        notes.append(
            "2枚の縦横比が少し違います。学習前にアプリ側で比較画像を元画像の比率へ合わせますが、"
            "できれば画像生成AIで元画像と同じ幅・高さの比較画像を作り直す方がきれいです。"
        )
    return notes


DEFAULT_REFERENCE_STYLE = "実写写真（強推奨）"


REFERENCE_STYLE_PRESETS = {
    "実写写真（強推奨）": {
        "description": (
            "完全な実写写真として再構成する。現実の人物を撮影したポートレートまたはコスプレ写真のようにし、"
            "元画像の線画、セル塗り、アニメ調、イラスト調、絵柄の癖は残さない。"
            "ただし外周シルエット、顔の輪郭、首、肩、服の輪郭は元画像に合わせる"
        ),
        "tags": [
            "photo_(medium)",
            "real_life",
            "photorealistic",
            "realistic",
            "portrait",
            "simple_background",
            "white_background",
            "natural_lighting",
            "detailed_skin",
            "detailed_hair",
            "fabric_texture",
        ],
        "negative_tags": [
            "anime",
            "manga",
            "illustration",
            "lineart",
            "cel_shading",
            "flat_color",
            "drawn_outline",
            "2d",
            "cartoon",
        ],
        "strict_instruction": (
            "このモードでは、かわいいアニメ風・厚塗りアニメ風・AIイラスト風に戻さないでください。"
            "線で描いた輪郭ではなく、実際の肌、髪、布、光でできた写真にしてください。"
            "ただし顔や体を現実の人間として自然に補正しないでください。"
            "元画像の外周シルエット、顔の横幅、頬、顎、首、肩、服の輪郭を最優先で保ちます。"
            "人体として不自然に見えても構いません。瞳のサイズ、虹彩、瞳孔、ハイライトは実写方向へ変えて構いません。"
        ),
    },
    "3Dレンダー": {
        "description": (
            "滑らかな3Dキャラクターレンダーとして再構成する。線画を消し、立体モデル、CGシェーダー、"
            "なめらかな髪と布の質感にする。人物や服、ポーズは変えない"
        ),
        "tags": [
            "3d",
            "cg",
            "3d_render",
            "simple_background",
            "white_background",
            "smooth_shading",
            "soft_lighting",
            "detailed",
        ],
        "negative_tags": [
            "lineart",
            "cel_shading",
            "flat_color",
            "sketch",
            "rough_drawing",
        ],
        "strict_instruction": (
            "このモードでは、元画像の線画や平面的なセル塗りを残さず、立体物としてレンダリングしてください。"
        ),
    },
    "全身・ちびキャラ用": {
        "description": (
            "完全な実写写真として再構成する。現実の人物を撮影したポートレートまたはコスプレ写真のようにし、"
            "元画像の線画、セル塗り、アニメ調、イラスト調、絵柄の癖は残さない。"
        ),
        "tags": [
            "full_body",
            "photo_(medium)",
            "real_life",
            "photorealistic",
            "realistic",
            "simple_background",
            "white_background",
            "natural_lighting",
            "detailed_skin",
            "detailed_hair",
            "fabric_texture",
        ],
        "negative_tags": [
            "chibi",
            "super_deformed",
            "anime",
            "manga",
            "illustration",
            "cg_illustration",
            "ai_illustration",
            "lineart",
            "cel_shading",
            "flat_color",
            "drawn_outline",
            "2d",
            "cartoon",
            "hand_drawn",
        ],
        "strict_instruction": (
            "このモードでは、かわいいアニメ風・厚塗りアニメ風・AIイラスト風に戻さないでください。"
            "線で描いた輪郭ではなく、実際の肌、髪、布、光でできた写真にしてください。"
            "ただし顔や体を現実の人間として自然に補正しないでください。"
            "瞳のサイズ、虹彩、瞳孔、ハイライトは実写方向へ変えて構いません。"
        ),
    },
    "厚塗り・油彩": {
        "description": (
            "厚塗りまたは油彩寄りのペインティングにする。筆致、面の重なり、絵具の質感を出し、"
            "元画像の細い線画とフラットなセル塗りは避ける"
        ),
        "tags": [
            "painting",
            "oil_painting",
            "painterly",
            "brush_stroke",
            "simple_background",
            "white_background",
            "soft_lighting",
        ],
        "negative_tags": [
            "clean_lineart",
            "cel_shading",
            "flat_color",
            "anime_screencap",
        ],
        "strict_instruction": (
            "このモードでは、線画の上に色を置いた絵ではなく、面と筆致で形を作る画像にしてください。"
        ),
    },
    "ニュートラルアニメ": {
        "description": "癖の少ない標準的なアニメ調にする。構図確認用として使い、元画像の固有画風は避ける",
        "tags": [
            "anime_style",
            "simple_background",
            "white_background",
            "clean_lineart",
            "flat_color",
            "simple_shading",
            "neutral_lighting",
        ],
        "negative_tags": [
            "artist_name",
            "signature",
            "watermark",
        ],
        "strict_instruction": (
            "このモードは差が弱くなりやすいので、元画像と似すぎる場合は実写写真または3Dレンダーで作り直してください。"
        ),
    },
    "低情報フラット": {
        "description": "情報量を抑えたフラットな塗りにする。形の確認用で、強い個性や凝った筆致を避ける",
        "tags": [
            "anime_style",
            "simple_background",
            "white_background",
            "flat_color",
            "minimal_shading",
            "clean_lineart",
        ],
        "negative_tags": [
            "painterly",
            "oil_painting",
            "photorealistic",
            "detailed_texture",
        ],
        "strict_instruction": (
            "このモードは元画像がすでにフラットな場合は差が弱くなります。その場合は実写写真か3Dレンダーを使ってください。"
        ),
    },
}

STYLE_NAME_ALIASES = {
    "実写寄り": "実写写真（強推奨）",
    "実写": "実写写真（強推奨）",
    "3D CG寄り": "3Dレンダー",
    "ちび": "全身・ちびキャラ用",
    "ちびキャラ": "全身・ちびキャラ用",
    "フラット塗り": "低情報フラット",
}


NEGATIVE_REFERENCE_TAGS = [
    "artist_name",
    "signature",
    "watermark",
    "text",
    "logo",
    "bad_hands",
    "extra_fingers",
    "different_pose",
    "different_clothes",
    "different_character",
    "cropped",
    "lowres",
    "blurry",
]


def style_preset(style_name: str | None) -> dict[str, list[str] | str]:
    normalized = STYLE_NAME_ALIASES.get(style_name or "", style_name or DEFAULT_REFERENCE_STYLE)
    return REFERENCE_STYLE_PRESETS.get(normalized, REFERENCE_STYLE_PRESETS[DEFAULT_REFERENCE_STYLE])


def build_pair_prompt(
    style_name: str | None = None,
    content_caption: str | None = None,
    image_size: tuple[int, int] | None = None,
) -> str:
    normalized_style_name = STYLE_NAME_ALIASES.get(style_name or "", style_name or DEFAULT_REFERENCE_STYLE)
    preset = style_preset(style_name)
    clean_content_caption = clean_caption_reference(content_caption)
    content_line = f"\n\n内容タグの参考:\n{clean_content_caption}" if clean_content_caption else ""
    size_line = ""
    if image_size:
        width, height = image_size
        aspect = width / max(height, 1)
        size_line = (
            f"- 出力サイズは元画像と同じ {width}x{height}、または完全に同じ縦横比 {aspect:.4f}:1 にする\n"
            "- 正方形に変えない。上下左右に余白を足さない。顔や体の位置、拡大率、カメラ距離も変えない\n"
        )
    else:
        size_line = (
            "- 出力サイズは元画像と同じ幅・高さ、または完全に同じ縦横比にする\n"
            "- 正方形に変えない。上下左右に余白を足さない。顔や体の位置、拡大率、カメラ距離も変えない\n"
        )
    if normalized_style_name == "全身・ちびキャラ用":
        return (
            "添付画像を参照して、人物・ポーズ・表情・服・髪型・構図・背景の位置関係をできるだけ同じに保ったまま、"
            "画風と頭身だけを大きく変えたLoRA学習用の比較画像を1枚作ってください。"
            f"{content_line}\n\n"
            f"比較画像の方向:\n{preset['description']}\n\n"
            f"重要:\n{preset.get('strict_instruction', '')}\n\n"
            "条件:\n"
            "- キャラクターや物体、服、ポーズ、カメラ距離、背景配置は変えない\n"
            "- 線の太さ、塗り、影、質感、色作り、描画媒体だけを別方向にする\n"
            "- 元画像の作者名、作品名、固有画風を真似する指定は入れない\n"
            "- ロゴ、文字、署名、透かし、余計な小物は入れない\n"
            f"{size_line}"
            "- 完成画像は、元画像と並べて見た時に内容が同じで、画風と頭身だけ違う状態にする\n\n"
            "失敗条件:\n"
            "- 元画像とほぼ同じ絵柄、線画、セル塗り、色作りが残っている\n"
            "- 実写写真モードなのに、アニメ風・イラスト風・CGイラスト風に見える\n"
            "- 顔や体型を現実の人間として自然に補正している\n"
            "- 顔の輪郭、首、肩、服の外側のシルエットが元画像から変わっている\n"
            "- 服、ポーズ、構図、トリミング、縦横比が元画像から変わっている"
        )
    return (
        "添付画像を参照して、人物・ポーズ・表情・服・髪型・構図・背景の位置関係をできるだけ同じに保ったまま、"
        "画風だけを大きく変えたLoRA学習用の比較画像を1枚作ってください。"
        f"{content_line}\n\n"
        f"比較画像の方向:\n{preset['description']}\n\n"
        f"重要:\n{preset.get('strict_instruction', '')}\n\n"
        "条件:\n"
        "- 最優先: 元画像の外周シルエットをできるだけ同じにする\n"
        "- 顔の横幅、頬の丸み、顎の形、頭の大きさ、首の長さ、肩幅、胴体、服の外側の輪郭を変えない\n"
        "- 人体として自然に見えるよう補正しない。現実の人間として不自然でも、元画像の輪郭を優先する\n"
        "- 瞳のサイズ、虹彩、瞳孔、ハイライト、目の質感は比較画像の方向に合わせて大きく変えてよい\n"
        "- キャラクターや物体、服、ポーズ、カメラ距離、背景配置は変えない\n"
        "- 線の太さ、塗り、影、質感、色作り、描画媒体だけを別方向にする\n"
        "- 元画像の作者名、作品名、固有画風を真似する指定は入れない\n"
        "- ロゴ、文字、署名、透かし、余計な小物は入れない\n"
        f"{size_line}"
        "- 完成画像は、元画像と並べて見た時に内容が同じで、画風だけ違う状態にする\n\n"
        "失敗条件:\n"
        "- 元画像とほぼ同じ絵柄、線画、セル塗り、色作りが残っている\n"
        "- 実写写真モードなのに、アニメ風・イラスト風・CGイラスト風に見える\n"
        "- 顔や体型を現実の人間として自然に補正している\n"
        "- 顔の輪郭、首、肩、服の外側のシルエットが元画像から変わっている\n"
        "- 服、ポーズ、構図、トリミング、縦横比が元画像から変わっている"
    )


def build_danbooru_prompt(tags: Iterable[str], style_name: str | None = None) -> str:
    preset = style_preset(style_name)
    content = content_tags(tags)[:28]
    style_tags = [clean_tag(tag) for tag in preset["tags"]]  # type: ignore[index]
    merged = dedupe(content + style_tags)
    return ", ".join(merged)


def build_negative_prompt(style_name: str | None = None) -> str:
    preset = style_preset(style_name)
    negative_tags = list(NEGATIVE_REFERENCE_TAGS) + list(preset.get("negative_tags", []))  # type: ignore[arg-type]
    return ", ".join(dedupe(negative_tags))


def build_reference_prompt_bundle(style_image: str | Path, style_name: str | None = None) -> tuple[str, str, str, str, str, str]:
    analysis = analyze_image(style_image)
    target_caption = analysis.caption
    reference_caption = target_caption
    image_size = (int(analysis.traits["width"]), int(analysis.traits["height"]))
    instruction = build_pair_prompt(style_name, target_caption, image_size)
    danbooru_prompt = build_danbooru_prompt(analysis.tags, style_name)
    negative_prompt = build_negative_prompt(style_name)
    notes = [
        f"学習画像: {analysis.method}",
        f"推定: {analysis.traits['color_mode']} / {analysis.traits['line_mode']} / {analysis.traits['width']}x{analysis.traits['height']}",
        f"比較画像の推奨サイズ: {analysis.traits['width']}x{analysis.traits['height']}（同じ縦横比を厳守）",
        "画像生成AIでは、下の指示文とdanbooruタグ案を使って比較画像を作ってください。",
        "比較画像ができたら、人物・服・ポーズ・構図が元画像と同じか確認してから読み込んでください。",
    ]
    notes.extend(analysis.notes)
    return target_caption, reference_caption, instruction, danbooru_prompt, negative_prompt, "\n".join(notes)


def auto_caption_pair(
    style_image: str | Path,
    pair_image: str | Path | None = None,
    detailed: bool = False,
) -> tuple[str, str, str, str]:
    general_threshold = DETAILED_GENERAL_TAG_THRESHOLD if detailed else DEFAULT_GENERAL_TAG_THRESHOLD
    max_tags = DETAILED_MAX_CONTENT_TAGS if detailed else DEFAULT_MAX_CONTENT_TAGS
    style = analyze_image(style_image, general_threshold=general_threshold, max_tags=max_tags)
    pair = analyze_image(pair_image, general_threshold=general_threshold, max_tags=max_tags) if pair_image else None

    if detailed and pair:
        shared_caption = make_caption(dedupe(style.tags + pair.tags), max_tags=max_tags)
    else:
        shared_caption = style.caption
    pair_caption = shared_caption if pair else ""

    notes = [
        f"学習画像: {style.method}",
        f"推定: {style.traits['color_mode']} / {style.traits['line_mode']} / {style.traits['width']}x{style.traits['height']}",
    ]
    if detailed:
        notes.append(
            f"詳細タグ付け: 一般タグしきい値 {general_threshold} / 最大 {max_tags} タグ"
            + (" / 比較画像側の内容タグも追加" if pair else "")
        )
    notes.extend(style.notes)
    if pair:
        notes.append(f"比較画像: {pair.method}")
        notes.append(f"推定: {pair.traits['color_mode']} / {pair.traits['line_mode']} / {pair.traits['width']}x{pair.traits['height']}")
        notes.extend(pair.notes)
    notes.extend(pair_similarity_notes(style, pair))

    advice = (
        "おすすめ: 2枚のキャプションは同じ内容タグにそろえ、画風タグは入れないでください。"
        "差分LoRAは「同じ内容で違う画風」の差を抜き出すほどきれいになります。"
    )
    image_size = (int(style.traits["width"]), int(style.traits["height"]))
    return shared_caption, pair_caption, "\n".join(notes + [advice]), build_pair_prompt(content_caption=shared_caption, image_size=image_size)
