"""AI generation helpers backed by configurable HTTP model endpoints."""

import base64
import copy
import json
import mimetypes
import os
import time
import urllib.error
import urllib.request
import uuid

from . import appconfig


class AIError(RuntimeError):
    """Raised when AI generation cannot complete."""


_RESOLUTION_SIZES = {
    "512": (512, 512),
    "1k": (1024, 1024),
    "2k": (2048, 2048),
    "4k": (4096, 4096),
}


def gen_title(
    title_prompt,
    old_title,
    retry=None,
    config=None,
    models_path=appconfig.AI_MODELS_PATH,
):
    """Generate a new product title from a prompt and the old title."""

    cfg = appconfig.load_config() if config is None else config
    ai_cfg = appconfig.ai_config(cfg)
    model = _role_model("text", ai_cfg.get("default_text_model"), models_path)
    payload = _chat_payload(
        model,
        [
            {"role": "system", "content": str(title_prompt or "").strip()},
            {
                "role": "user",
                "content": "旧标题：\n%s\n\n请只返回新标题，不要解释。" % str(old_title or ""),
            },
        ],
    )
    attempts = _attempt_count(ai_cfg, retry)
    data = _call_with_retry(
        model,
        payload,
        cfg,
        attempts,
        request_kind="json",
    )
    text = _extract_text(data).strip()
    if not text:
        raise AIError("AI 返回为空标题")
    return text


def gen_cover(
    cover_prompt,
    old_cover_path,
    out_path,
    resolution=None,
    jpg_quality=None,
    retry=None,
    config=None,
    models_path=appconfig.AI_MODELS_PATH,
):
    """Generate a new cover image and save it as a JPEG file."""

    old_cover_path = os.path.abspath(str(old_cover_path))
    if not os.path.exists(old_cover_path):
        raise FileNotFoundError("旧封面图片不存在: %s" % old_cover_path)
    if not out_path:
        raise AIError("缺少新封面输出路径")

    cfg = appconfig.load_config() if config is None else config
    ai_cfg = appconfig.ai_config(cfg)
    model = _role_model("image", ai_cfg.get("default_image_model"), models_path)
    resolution = str(resolution or ai_cfg.get("resolution", "1k"))
    quality = _jpg_quality(jpg_quality if jpg_quality is not None else ai_cfg.get("jpg_quality", 90))
    attempts = _attempt_count(ai_cfg, retry)

    api_type = model.get("api_type", "auto")
    if api_type == "images_edits":
        body, content_type = _image_edit_body(model, cover_prompt, old_cover_path, resolution)
        data = _call_with_retry(
            model,
            body,
            cfg,
            attempts,
            request_kind="multipart",
            content_type=content_type,
        )
    else:
        payload = _image_chat_payload(model, cover_prompt, old_cover_path, resolution)
        data = _call_with_retry(model, payload, cfg, attempts, request_kind="json")

    image_bytes = _extract_image_bytes(data, model, cfg)
    return _save_jpeg(image_bytes, out_path, resolution, quality)


def _role_model(category, name, models_path):
    if not name:
        raise AIError("未配置默认 %s 模型" % category)
    model = appconfig.get_model(name, path=models_path)
    if model.get("category") != category:
        raise AIError("模型 %s 不是 %s 类别" % (name, category))
    if not model.get("enabled", True):
        raise AIError("模型已禁用: %s" % name)
    missing = [
        field
        for field in ("url", "model", "api_key")
        if not str(model.get(field, "")).strip()
    ]
    if missing:
        raise AIError("模型 %s 缺少字段: %s" % (name, ", ".join(missing)))
    return model


def _attempt_count(ai_cfg, retry):
    retry_count = ai_cfg.get("retry", 2) if retry is None else retry
    return max(1, int(retry_count) + 1)


def _read_timeout(model, config):
    return int(model.get("timeout_seconds") or appconfig.response_timeout(config))


def _connect_timeout(model):
    return int(model.get("connect_timeout_seconds") or 30)


def _headers(model, content_type):
    return {
        "Authorization": "Bearer " + model["api_key"],
        "Content-Type": content_type,
    }


def _call_with_retry(model, body, config, attempts, request_kind, content_type=None):
    last_exc = None
    for index in range(attempts):
        try:
            return _call_once(model, body, config, request_kind, content_type=content_type)
        except Exception as exc:
            last_exc = exc
            if index + 1 >= attempts:
                break
            time.sleep(min(2.0, 0.4 * (index + 1)))
    raise AIError(
        "AI 调用失败（已尝试 %s 次）: %s"
        % (attempts, _redact(str(last_exc), model))
    ) from last_exc


def _call_once(model, body, config, request_kind, content_type=None):
    if request_kind == "json":
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
        content_type = "application/json"
    else:
        data = body
    request = urllib.request.Request(
        model["url"],
        data=data,
        headers=_headers(model, content_type),
        method="POST",
    )
    timeout = max(_connect_timeout(model), _read_timeout(model, config))
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            raw = response.read()
    except urllib.error.HTTPError as exc:
        detail = _safe_http_error(exc, model)
        raise AIError("HTTP %s: %s" % (exc.code, detail)) from exc
    except urllib.error.URLError as exc:
        raise AIError(_redact(str(exc.reason), model)) from exc
    try:
        return json.loads(raw.decode("utf-8"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise AIError("AI 返回不是有效 JSON") from exc


def _safe_http_error(exc, model):
    try:
        body = exc.read(2048).decode("utf-8", errors="replace")
    except Exception:
        body = ""
    return _redact(body or str(exc), model)


def _chat_payload(model, messages):
    payload = {
        "model": model["model"],
        "messages": messages,
    }
    payload.update(copy.deepcopy(model.get("extra_body", {})))
    return payload


def _image_chat_payload(model, cover_prompt, old_cover_path, resolution):
    prompt = "%s\n\n目标分辨率：%s。" % (str(cover_prompt or "").strip(), resolution)
    content = [
        {"type": "text", "text": prompt.strip()},
        {
            "type": "image_url",
            "image_url": {"url": _image_data_url(old_cover_path)},
        },
    ]
    return _chat_payload(model, [{"role": "user", "content": content}])


def _image_edit_body(model, cover_prompt, old_cover_path, resolution):
    fields = {
        "model": model["model"],
        "prompt": str(cover_prompt or ""),
        "size": _resolution_size_text(resolution),
    }
    fields.update(copy.deepcopy(model.get("extra_body", {})))
    files = {
        "image": (
            os.path.basename(old_cover_path),
            open(old_cover_path, "rb").read(),
            mimetypes.guess_type(old_cover_path)[0] or "application/octet-stream",
        )
    }
    return _multipart_body(fields, files)


def _multipart_body(fields, files):
    boundary = "----cmshopee-%s" % uuid.uuid4().hex
    chunks = []
    for name, value in fields.items():
        chunks.extend(
            [
                ("--%s\r\n" % boundary).encode("utf-8"),
                ('Content-Disposition: form-data; name="%s"\r\n\r\n' % name).encode("utf-8"),
                str(value).encode("utf-8"),
                b"\r\n",
            ]
        )
    for name, file_info in files.items():
        filename, data, content_type = file_info
        chunks.extend(
            [
                ("--%s\r\n" % boundary).encode("utf-8"),
                (
                    'Content-Disposition: form-data; name="%s"; filename="%s"\r\n'
                    % (name, filename)
                ).encode("utf-8"),
                ("Content-Type: %s\r\n\r\n" % content_type).encode("utf-8"),
                data,
                b"\r\n",
            ]
        )
    chunks.append(("--%s--\r\n" % boundary).encode("utf-8"))
    return b"".join(chunks), "multipart/form-data; boundary=%s" % boundary


def _image_data_url(path):
    mime = mimetypes.guess_type(path)[0] or "image/jpeg"
    with open(path, "rb") as fh:
        encoded = base64.b64encode(fh.read()).decode("ascii")
    return "data:%s;base64,%s" % (mime, encoded)


def _extract_text(data):
    if isinstance(data, dict):
        for key in ("output_text", "text", "content"):
            value = data.get(key)
            if isinstance(value, str):
                return value
        choices = data.get("choices")
        if isinstance(choices, list) and choices:
            first = choices[0]
            if isinstance(first, dict):
                if isinstance(first.get("text"), str):
                    return first["text"]
                message = first.get("message") or {}
                content = message.get("content")
                return _content_text(content)
    return ""


def _content_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if not isinstance(item, dict):
                continue
            value = item.get("text") or item.get("content")
            if isinstance(value, str):
                parts.append(value)
        return "".join(parts)
    return ""


def _extract_image_bytes(data, model, config):
    image_ref = _find_image_ref(data)
    if not image_ref:
        raise AIError("AI 返回中没有图片数据")
    if image_ref.startswith("data:"):
        return _decode_data_url(image_ref)
    if _looks_base64(image_ref):
        return base64.b64decode(image_ref)
    return _download_image(image_ref, model, config)


def _find_image_ref(value):
    if isinstance(value, dict):
        for key in ("b64_json", "base64", "image_base64", "image", "url"):
            candidate = value.get(key)
            if isinstance(candidate, str) and candidate.strip():
                return candidate.strip()
        image_url = value.get("image_url")
        if isinstance(image_url, str):
            return image_url
        if isinstance(image_url, dict):
            candidate = image_url.get("url")
            if isinstance(candidate, str):
                return candidate
        for key in ("data", "choices", "output", "content", "images"):
            candidate = _find_image_ref(value.get(key))
            if candidate:
                return candidate
        message = value.get("message")
        if message is not None:
            candidate = _find_image_ref(message)
            if candidate:
                return candidate
    elif isinstance(value, list):
        for item in value:
            candidate = _find_image_ref(item)
            if candidate:
                return candidate
    return None


def _decode_data_url(value):
    if "," not in value:
        raise AIError("图片 data URL 格式错误")
    return base64.b64decode(value.split(",", 1)[1])


def _looks_base64(value):
    compact = value.strip()
    if compact.startswith(("http://", "https://")):
        return False
    if len(compact) < 32:
        return False
    allowed = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=\n\r")
    return all(char in allowed for char in compact)


def _download_image(url, model, config):
    request = urllib.request.Request(url, method="GET")
    try:
        with urllib.request.urlopen(request, timeout=_read_timeout(model, config)) as response:
            return response.read()
    except Exception as exc:
        raise AIError("下载 AI 图片失败: %s" % _redact(str(exc), model)) from exc


def _save_jpeg(image_bytes, out_path, resolution, jpg_quality):
    try:
        from PIL import Image
    except ImportError as exc:
        raise AIError("缺少 Pillow，无法保存 AI 图片") from exc

    import io

    out_path = os.path.abspath(str(out_path))
    directory = os.path.dirname(out_path)
    if directory:
        os.makedirs(directory, exist_ok=True)
    size = _resolution_size(resolution)
    try:
        with Image.open(io.BytesIO(image_bytes)) as image:
            image = image.convert("RGB")
            if size:
                image = image.resize(size, Image.LANCZOS)
            image.save(out_path, "JPEG", quality=jpg_quality, optimize=True)
    except Exception as exc:
        raise AIError("AI 图片保存失败: %s" % exc) from exc
    return out_path


def _resolution_size(resolution):
    return _RESOLUTION_SIZES.get(str(resolution))


def _resolution_size_text(resolution):
    size = _resolution_size(resolution)
    if not size:
        return str(resolution)
    return "%sx%s" % size


def _jpg_quality(value):
    value = int(value)
    return min(100, max(1, value))


def _redact(text, model):
    result = str(text)
    for secret in (model.get("api_key"),):
        if secret:
            result = result.replace(str(secret), "***")
    return result