"""Safe local export helpers for AI image studio final selections."""

from __future__ import annotations

import os
import re
import shutil
import uuid
from dataclasses import dataclass
from datetime import datetime

from . import appconfig, image_studio


EXISTING_FAIL = "fail"
EXISTING_OVERWRITE_MANAGED = "overwrite_managed"
EXISTING_TIMESTAMP = "timestamp"
EXISTING_MODES = {EXISTING_FAIL, EXISTING_OVERWRITE_MANAGED, EXISTING_TIMESTAMP}


class ImageStudioExportError(RuntimeError):
    """Raised when AI studio selections cannot be exported safely."""


class ExportTargetExistsError(ImageStudioExportError):
    """Raised when the default target directory already exists."""


@dataclass(frozen=True)
class ExportedFile:
    selection_type: str
    asset_id: int
    source_path: str
    output_path: str


@dataclass(frozen=True)
class ExportResult:
    target_dir: str
    files: tuple[ExportedFile, ...]
    main_count: int
    detail_count: int
    existing_mode: str


@dataclass(frozen=True)
class GenerationRoundExportedFile:
    job_id: int
    asset_id: int
    job_type: str
    source_path: str
    output_path: str


@dataclass(frozen=True)
class GenerationRoundExportResult:
    target_dir: str
    files: tuple[GenerationRoundExportedFile, ...]
    skipped_count: int
    cancelled: bool


def target_dir_for_project(project, parent_dir, suffix=None):
    parent = _existing_parent_dir(parent_dir)
    item = _safe_item_id(getattr(project, "item_id", ""))
    dirname = item if not suffix else f"{item}_{suffix}"
    return os.path.abspath(os.path.join(parent, dirname))


def export_project_selection(
    project_id,
    parent_dir,
    *,
    existing_mode=EXISTING_FAIL,
    path=None,
    config=None,
    timestamp=None,
):
    """Export current main/detail selections as ordered JPEG files."""

    mode = str(existing_mode or EXISTING_FAIL)
    if mode not in EXISTING_MODES:
        raise ImageStudioExportError("导出目录处理方式无效")
    cfg = appconfig.load_config() if config is None else config
    database_path = path or appconfig.db_path(cfg)
    project = image_studio.get_project(project_id, path=database_path)
    if project is None:
        raise ImageStudioExportError("AI工场项目不存在")
    parent = _existing_parent_dir(parent_dir)
    planned = _planned_files(project, database_path)
    if not planned:
        raise ImageStudioExportError("主图和详情图终选都为空，不能导出")
    target = _choose_target_dir(project, parent, mode, timestamp=timestamp)
    quality = int(appconfig.ai_config(cfg).get("jpg_quality", 90) or 90)
    staging = os.path.join(parent, f".{_safe_item_id(project.item_id)}_staging_{uuid.uuid4().hex}")
    try:
        os.makedirs(staging, exist_ok=False)
        staged = []
        for item in planned:
            output_path = os.path.join(staging, item["filename"])
            _save_jpeg(item["source_path"], output_path, quality)
            staged.append(ExportedFile(item["selection_type"], item["asset_id"], item["source_path"], output_path))
        os.makedirs(target, exist_ok=True)
        if os.path.exists(target) and mode == EXISTING_OVERWRITE_MANAGED:
            _remove_managed_exports(target, project.item_id)
        final_files = []
        for staged_file in staged:
            final_path = os.path.join(target, os.path.basename(staged_file.output_path))
            os.replace(staged_file.output_path, final_path)
            final_files.append(
                ExportedFile(
                    staged_file.selection_type,
                    staged_file.asset_id,
                    staged_file.source_path,
                    final_path,
                )
            )
        return ExportResult(
            target_dir=target,
            files=tuple(final_files),
            main_count=sum(1 for item in final_files if item.selection_type == "main"),
            detail_count=sum(1 for item in final_files if item.selection_type == "detail"),
            existing_mode=mode,
        )
    except Exception as exc:
        if isinstance(exc, ImageStudioExportError):
            raise
        raise ImageStudioExportError(f"导出终选图片失败：{exc}") from exc
    finally:
        if os.path.isdir(staging):
            shutil.rmtree(staging, ignore_errors=True)


def export_generation_round(
    project_id,
    generation_round_key,
    parent_dir,
    *,
    path=None,
    should_stop=None,
):
    """Copy usable successful output images from one persisted generation round."""

    project = image_studio.get_project(project_id, path=path)
    if project is None:
        raise ImageStudioExportError("商品套图项目不存在，无法导出历史记录")
    try:
        parent = _existing_parent_dir(parent_dir)
    except ImageStudioExportError as exc:
        raise ImageStudioExportError("导出父目录不可用") from exc
    jobs = image_studio.list_generation_round_current_jobs(
        project.id,
        generation_round_key,
        path=path,
    )
    planned, skipped_count = _planned_generation_round_files(jobs, path)
    if not planned:
        raise ImageStudioExportError("本轮没有可导出的成功生成图片")

    target = _create_generation_round_target_dir(project, parent, jobs)
    copied = []
    cancelled = False
    try:
        for item in planned:
            if callable(should_stop) and should_stop():
                cancelled = True
                break
            output_path = os.path.join(target, item["filename"])
            try:
                shutil.copy2(item["source_path"], output_path)
            except OSError:
                skipped_count += 1
                if os.path.isfile(output_path):
                    try:
                        os.remove(output_path)
                    except OSError:
                        pass
                continue
            copied.append(
                GenerationRoundExportedFile(
                    job_id=int(item["job_id"]),
                    asset_id=int(item["asset_id"]),
                    job_type=str(item["job_type"]),
                    source_path=item["source_path"],
                    output_path=output_path,
                )
            )
    except Exception as exc:
        raise ImageStudioExportError("导出本轮生成图片失败") from exc

    if not copied:
        try:
            os.rmdir(target)
        except OSError:
            pass
        if cancelled:
            return GenerationRoundExportResult("", (), skipped_count, True)
        raise ImageStudioExportError("本轮生成图片导出失败")
    return GenerationRoundExportResult(
        target_dir=target,
        files=tuple(copied),
        skipped_count=skipped_count,
        cancelled=cancelled,
    )


def _choose_target_dir(project, parent_dir, existing_mode, timestamp=None):
    target = target_dir_for_project(project, parent_dir)
    if not os.path.exists(target):
        return target
    if not os.path.isdir(target):
        raise ImageStudioExportError(f"导出目标已存在但不是目录：{target}")
    if existing_mode == EXISTING_FAIL:
        raise ExportTargetExistsError(f"商品目录已存在：{target}")
    if existing_mode == EXISTING_OVERWRITE_MANAGED:
        return target
    stamp = timestamp or datetime.now().strftime("%Y%m%d_%H%M%S")
    base = target_dir_for_project(project, parent_dir, suffix=stamp)
    candidate = base
    index = 2
    while os.path.exists(candidate):
        candidate = f"{base}_{index}"
        index += 1
    return candidate


def _planned_files(project, db_path):
    item = _safe_item_id(project.item_id)
    planned = []
    for selection_type, label in (("main", "主图"), ("detail", "详情图")):
        selections = image_studio.list_selections(project.id, selection_type, path=db_path)
        for index, selection in enumerate(selections, start=1):
            asset = image_studio.get_asset(selection.asset_id, path=db_path)
            if asset is None:
                raise ImageStudioExportError(f"终选照片不存在：#{selection.asset_id}")
            source_path = str(asset.local_path or "")
            if not source_path or not os.path.isfile(source_path):
                raise ImageStudioExportError(f"终选照片本地文件缺失：#{asset.id}")
            planned.append(
                {
                    "selection_type": selection_type,
                    "asset_id": int(asset.id),
                    "source_path": os.path.abspath(source_path),
                    "filename": f"{item}_{label}_{index}.jpg",
                }
            )
    return planned


def _planned_generation_round_files(jobs, db_path):
    planned = []
    skipped_count = 0
    for index, job in enumerate(jobs, 1):
        if str(getattr(job, "status", "") or "") != "succeeded":
            continue
        asset_id = getattr(job, "output_asset_id", None)
        asset = image_studio.get_asset(asset_id, path=db_path) if asset_id else None
        source_path = str(getattr(asset, "local_path", "") or "")
        if (
            asset is None
            or str(getattr(asset, "status", "") or "")
            == image_studio.ASSET_STATUS_MISSING
            or not source_path
            or not os.path.isfile(source_path)
        ):
            skipped_count += 1
            continue
        suffix = os.path.splitext(source_path)[1].lower()
        if not suffix or len(suffix) > 8:
            suffix = ".jpg"
        job_type = _safe_export_component(
            getattr(job, "job_type", ""),
            fallback="套图",
        )
        planned.append(
            {
                "job_id": int(job.id),
                "asset_id": int(asset.id),
                "job_type": str(getattr(job, "job_type", "") or "套图"),
                "source_path": os.path.abspath(source_path),
                "filename": "%02d_%s%s" % (index, job_type, suffix),
            }
        )
    return planned, skipped_count


def _create_generation_round_target_dir(project, parent_dir, jobs):
    account = _safe_export_component(
        getattr(project, "account_name", "") or getattr(project, "account_alias", ""),
        fallback="未命名店铺",
    )
    item = _safe_export_component(getattr(project, "item_id", ""), fallback="临时草稿")
    created_at = str(getattr(jobs[0], "created_at", "") or "") if jobs else ""
    digits = "".join(character for character in created_at if character.isdigit())[:14]
    stamp = digits or datetime.now().strftime("%Y%m%d_%H%M%S")
    base = os.path.join(parent_dir, "%s_%s_%s" % (account, item, stamp))
    candidate = base
    index = 2
    while True:
        try:
            os.makedirs(candidate, exist_ok=False)
            return candidate
        except FileExistsError:
            candidate = "%s_%d" % (base, index)
            index += 1
        except OSError as exc:
            raise ImageStudioExportError("无法创建导出目录") from exc


def _save_jpeg(source_path, output_path, quality):
    try:
        from PIL import Image

        with Image.open(source_path) as image:
            if image.mode in {"RGBA", "LA"} or (
                image.mode == "P" and "transparency" in getattr(image, "info", {})
            ):
                rgba = image.convert("RGBA")
                background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
                background.alpha_composite(rgba)
                final_image = background.convert("RGB")
            else:
                final_image = image.convert("RGB")
            final_image.save(output_path, format="JPEG", quality=max(1, min(95, int(quality or 90))))
        with Image.open(output_path) as check:
            check.verify()
    except Exception as exc:
        raise ImageStudioExportError(f"JPEG 转码失败：{exc}") from exc


def _remove_managed_exports(target_dir, item_id):
    item = re.escape(_safe_item_id(item_id))
    pattern = re.compile(rf"^{item}_(主图|详情图)_\d+\.jpg$", re.IGNORECASE)
    for filename in os.listdir(target_dir):
        if not pattern.match(filename):
            continue
        path = os.path.join(target_dir, filename)
        if os.path.isfile(path):
            os.remove(path)


def _existing_parent_dir(parent_dir):
    parent = os.path.abspath(str(parent_dir or ""))
    if not os.path.isdir(parent):
        raise ImageStudioExportError(f"导出父目录不存在：{parent}")
    return parent


def _safe_item_id(item_id):
    value = str(item_id or "").strip()
    if not value:
        raise ImageStudioExportError("商品ID不能为空")
    if any(char in value for char in ('/', '\\', os.sep, os.altsep or "\0")):
        raise ImageStudioExportError("商品ID包含非法路径字符")
    if value in {".", ".."}:
        raise ImageStudioExportError("商品ID不能是路径符号")
    safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in value).strip("_")
    if not safe:
        raise ImageStudioExportError("商品ID不能作为目录名")
    return safe


def _safe_export_component(value, *, fallback):
    text = str(value or "").strip()
    safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_")
    if safe:
        return safe
    return str(fallback)