from __future__ import annotations

import base64
import json
import mimetypes
import os
import traceback
import time
from dataclasses import dataclass
from datetime import datetime
from io import BytesIO
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import requests
from PIL import Image, ImageFile


DEFAULT_PROMPT = "高端电子产品海报，干净背景，商业摄影，细节清晰"
DEFAULT_MODEL = "gpt-image-2"
DEFAULT_PIXEL_SIZE = "4096x4096"
DEFAULT_QUALITY = "high"
DEFAULT_OUTPUT_FILE = "output.png"
DEFAULT_RESPONSE_DIR = "outputs/image-demo-runs"
DEFAULT_CONNECT_TIMEOUT = 10
DEFAULT_READ_TIMEOUT = 900
DEFAULT_IMAGE_FETCH_TIMEOUT = 120
DEFAULT_POLL_CONNECT_TIMEOUT = 10
DEFAULT_POLL_READ_TIMEOUT = 60
DEFAULT_POLL_INTERVAL = 5
DEFAULT_INITIAL_POLL_DELAY = 10
DEFAULT_MAX_POLL_WAIT = 1800
DIRECT_HOSTS = {"vip.auto-code.net","api.toskaxy.xyz","api.aoixx.com"}
SUPPORTED_ASPECT_RATIOS = {
    "1:1",
    "16:9",
    "9:16",
    "4:3",
    "3:4",
    "3:2",
    "2:3",
    "5:4",
    "4:5",
    "2:1",
    "1:2",
    "21:9",
    "9:21",
}
OPENAI_RATIO_SIZE_MAP = {
    "1:1": "1024x1024",
    "16:9": "1536x864",
    "9:16": "864x1536",
    "4:3": "1536x1152",
    "3:4": "1152x1536",
    "3:2": "1536x1024",
    "2:3": "1024x1536",
    "5:4": "1280x1024",
    "4:5": "1024x1280",
    "2:1": "1536x768",
    "1:2": "768x1536",
    "21:9": "1680x720",
    "9:21": "720x1680",
}


@dataclass(frozen=True)
class ImageInput:
    filename: str
    mime_type: str
    data: bytes


def load_env_file(path: str = ".env") -> None:
    env_path = Path(path)
    if not env_path.exists():
        return

    for raw_line in env_path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue

        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = value


def is_pixel_size(size: str) -> bool:
    parts = size.lower().split("x")
    return len(parts) == 2 and all(part.isdigit() for part in parts)


def normalize_size(size: str) -> str:
    value = size.strip().lower().replace(" ", "")
    if is_pixel_size(value):
        width, height = value.split("x", 1)
        return f"{int(width)}x{int(height)}"
    return value


def is_solid_black_image(img: Image.Image) -> bool:
    rgb = img.convert("RGB")
    return all(lo == 0 and hi == 0 for lo, hi in rgb.getextrema())


def env_int(name: str, default: int) -> int:
    raw_value = os.getenv(name, "").strip()
    if not raw_value:
        return default

    try:
        return int(raw_value)
    except ValueError as exc:
        raise RuntimeError(f"{name} 不是有效整数: {raw_value}") from exc


def make_run_dir() -> Path:
    response_root = Path(os.getenv("IMAGE_RESPONSE_DIR", DEFAULT_RESPONSE_DIR)).expanduser()
    run_id = datetime.now().strftime("%Y%m%d-%H%M%S-%f")
    run_dir = response_root / run_id
    run_dir.mkdir(parents=True, exist_ok=False)
    return run_dir


def dump_json(path: Path, data: Any) -> None:
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")


def save_http_response(path: Path, response: requests.Response) -> dict[str, Any]:
    record: dict[str, Any] = {
        "url": response.url,
        "status_code": response.status_code,
        "headers": dict(response.headers),
    }

    try:
        record["json"] = response.json()
    except ValueError:
        record["body_text"] = response.text

    dump_json(path, record)
    return record


def save_error_artifact(path: Path, error: BaseException, extra: dict[str, Any] | None = None) -> None:
    payload: dict[str, Any] = {
        "error_type": type(error).__name__,
        "error_message": str(error),
        "traceback": traceback.format_exc(),
    }
    if extra:
        payload.update(extra)
    dump_json(path, payload)


def resolve_output_path(run_dir: Path, output_file: str) -> Path:
    output_path = Path(output_file).expanduser()
    if output_path.is_absolute():
        return output_path
    return run_dir / output_path


def normalize_base_url(base_url: str) -> str:
    normalized = base_url.rstrip("/")
    if normalized.endswith("/v1"):
        return normalized[:-3]
    return normalized


def build_request_session(base_url: str) -> requests.Session:
    session = requests.Session()
    hostname = urlparse(base_url).hostname or ""
    if hostname in DIRECT_HOSTS:
        session.trust_env = False
        print(f"直连 {hostname}，忽略本地代理环境变量")
    return session


def resolve_request_style(base_url: str) -> str:
    request_style = os.getenv("IMAGE_REQUEST_STYLE", "auto").strip().lower()
    if request_style in {"openai", "apicodex"}:
        return request_style

    hostname = (urlparse(base_url).hostname or "").lower()
    if hostname.endswith("openai.com"):
        return "openai"
    return "apicodex"


def resolve_api_size(size: str, request_style: str) -> str:
    if request_style == "openai" and size in OPENAI_RATIO_SIZE_MAP:
        return OPENAI_RATIO_SIZE_MAP[size]
    return size


def resolve_image_source(prefix: str) -> str | None:
    candidates = {
        prefix: os.getenv(prefix, "").strip(),
        f"{prefix}_FILE": os.getenv(f"{prefix}_FILE", "").strip(),
        f"{prefix}_URL": os.getenv(f"{prefix}_URL", "").strip(),
        f"{prefix}_DATA_URL": os.getenv(f"{prefix}_DATA_URL", "").strip(),
    }
    provided = [(name, value) for name, value in candidates.items() if value]
    if len(provided) > 1:
        raise RuntimeError(
            f"{prefix} 相关变量只能设置一个，当前同时设置了: "
            + ", ".join(name for name, _ in provided)
        )
    if not provided:
        return None
    return provided[0][1]


def safe_filename(filename: str, fallback: str, mime_type: str) -> str:
    candidate = Path(filename).name.strip()
    if not candidate:
        candidate = fallback
    if not Path(candidate).suffix:
        candidate += mimetypes.guess_extension(mime_type) or ".png"
    return candidate


def data_url_to_image_input(data_url: str, fallback: str) -> ImageInput:
    if not data_url.startswith("data:image/") or "," not in data_url:
        raise RuntimeError("data URL 参考图格式无效")

    header, data_b64 = data_url.split(",", 1)
    if ";base64" not in header.lower():
        raise RuntimeError("data URL 参考图必须使用 base64 编码")

    mime_type = header[5:].split(";", 1)[0] or "image/png"
    try:
        data = base64.b64decode(data_b64, validate=True)
    except ValueError as exc:
        raise RuntimeError("data URL 参考图 base64 无法解码") from exc

    return ImageInput(
        filename=safe_filename("", fallback, mime_type),
        mime_type=mime_type,
        data=data,
    )


def local_file_to_image_input(file_path: str, fallback: str) -> ImageInput:
    path = Path(file_path).expanduser()
    if not path.is_absolute():
        path = Path.cwd() / path
    if not path.exists():
        raise FileNotFoundError(f"找不到参考图文件: {file_path}")

    mime_type = mimetypes.guess_type(path.name)[0] or "image/png"
    return ImageInput(
        filename=safe_filename(path.name, fallback, mime_type),
        mime_type=mime_type,
        data=path.read_bytes(),
    )


def fetch_url_to_image_input(image_url: str, fallback: str, session: requests.Session) -> ImageInput:
    resp = session.get(image_url, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_IMAGE_FETCH_TIMEOUT))
    resp.raise_for_status()

    mime_type = (resp.headers.get("Content-Type") or "").split(";", 1)[0].strip()
    if not mime_type.startswith("image/"):
        mime_type = mimetypes.guess_type(urlparse(image_url).path)[0] or "image/png"

    filename = safe_filename(Path(urlparse(image_url).path).name, fallback, mime_type)
    return ImageInput(filename=filename, mime_type=mime_type, data=resp.content)


def image_source_to_input(source: str, fallback: str, session: requests.Session) -> ImageInput:
    value = source.strip()
    lowered = value.lower()
    if lowered.startswith("data:image/"):
        return data_url_to_image_input(value, fallback)
    if lowered.startswith(("http://", "https://")):
        return fetch_url_to_image_input(value, fallback, session)
    return local_file_to_image_input(value, fallback)


def build_payload(base_url: str) -> tuple[dict[str, Any], str, str, str, str, str | None, str | None]:
    model = os.getenv("IMAGE_MODEL", DEFAULT_MODEL).strip()
    prompt = os.getenv("IMAGE_PROMPT", DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
    requested_size = normalize_size(os.getenv("IMAGE_SIZE", DEFAULT_PIXEL_SIZE))
    resolution = os.getenv("IMAGE_RESOLUTION", "").strip()
    quality = os.getenv("IMAGE_QUALITY", DEFAULT_QUALITY).strip()
    output_file = os.getenv("OUTPUT_FILE", DEFAULT_OUTPUT_FILE).strip()
    request_style = resolve_request_style(base_url)
    if not (is_pixel_size(requested_size) or requested_size in SUPPORTED_ASPECT_RATIOS):
        raise RuntimeError(f"不支持的 IMAGE_SIZE: {requested_size}")
    api_size = resolve_api_size(requested_size, request_style)

    payload: dict[str, Any] = {
        "model": model,
        "prompt": prompt,
        "n": int(os.getenv("IMAGE_N", "1")),
        "size": api_size,
        "quality": quality,
    }
    if not (request_style == "openai" and model.startswith("gpt-image")):
        payload["response_format"] = "b64_json"

    if not is_pixel_size(requested_size) and request_style != "openai":
        payload["resolution"] = resolution or "4K"

    reference_source = resolve_image_source("REFERENCE_IMAGE")
    mask_source = resolve_image_source("MASK_IMAGE")
    if mask_source and not reference_source:
        raise RuntimeError("MASK_IMAGE 只能和 REFERENCE_IMAGE 一起使用")

    request_mode = "edit" if reference_source else "generate"
    return payload, model, requested_size, output_file, request_mode, reference_source, mask_source


def payload_to_form_fields(payload: dict[str, Any]) -> dict[str, str]:
    fields: dict[str, str] = {}
    for key, value in payload.items():
        if value is None:
            continue
        fields[key] = str(value)
    return fields


def build_edit_files(
    request_style: str,
    reference_input: ImageInput,
    mask_input: ImageInput | None,
) -> list[tuple[str, tuple[str, bytes, str]]]:
    image_field_name = "image" if request_style == "openai" else "image[]"
    files: list[tuple[str, tuple[str, bytes, str]]] = [
        (
            image_field_name,
            (reference_input.filename, reference_input.data, reference_input.mime_type),
        )
    ]
    if mask_input is not None:
        files.append(("mask", (mask_input.filename, mask_input.data, mask_input.mime_type)))
    return files


def resolve_min_output_edge(requested_size: str, resolution: str | None, request_style: str) -> int | None:
    if is_pixel_size(requested_size):
        width, height = requested_size.split("x", 1)
        min_edge = max(int(width), int(height))
        return min_edge if min_edge >= 3840 else None
    if resolution and resolution.strip().upper() == "4K" and request_style != "openai":
        return 3840
    return None


def save_image_bytes(image_bytes: bytes, output_path: Path, min_output_edge: int | None = None) -> tuple[int, int]:
    output_path.parent.mkdir(parents=True, exist_ok=True)

    previous_truncated_setting = ImageFile.LOAD_TRUNCATED_IMAGES
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    try:
        with Image.open(BytesIO(image_bytes)) as img:
            img.load()
            width, height = img.size
            if is_solid_black_image(img):
                raise RuntimeError(
                    "上游返回了疑似纯黑占位图 "
                    f"({width}x{height})，通常表示请求没有走到正常生成链路。"
                    "请检查 OPENAI_BASE_URL、IMAGE_REQUEST_STYLE、model 和 size，"
                    "必要时切换到 vip.auto-code.net。"
                )

            if min_output_edge is not None and max(width, height) < min_output_edge:
                raise RuntimeError(f"生成结果未达到预期尺寸下限 {min_output_edge}px：{width}x{height}")
    finally:
        ImageFile.LOAD_TRUNCATED_IMAGES = previous_truncated_setting

    output_path.write_bytes(image_bytes)
    print(f"已保存 {output_path} ({width}x{height})")

    return width, height


def extract_image_bytes_from_item(item: dict[str, Any], session: requests.Session) -> bytes:
    if item.get("b64_json"):
        return base64.b64decode(item["b64_json"])

    url_value = item.get("url")
    if isinstance(url_value, list) and url_value:
        image_url = url_value[0]
    elif isinstance(url_value, str):
        image_url = url_value
    else:
        image_url = None

    if image_url:
        resp = session.get(image_url, timeout=(DEFAULT_CONNECT_TIMEOUT, DEFAULT_IMAGE_FETCH_TIMEOUT))
        resp.raise_for_status()
        return resp.content

    raise RuntimeError(f"响应里没有可用的图片数据: {item}")


def poll_task(
    task_id: str,
    base_url: str,
    headers: dict[str, str],
    session: requests.Session,
    run_dir: Path,
    max_wait_seconds: int,
    initial_poll_delay: int,
    poll_interval: int,
) -> bytes:
    task_url = f"{base_url}/v1/tasks/{task_id}"
    deadline = time.monotonic() + max_wait_seconds
    poll_index = 0

    time.sleep(initial_poll_delay)
    while True:
        if time.monotonic() >= deadline:
            raise TimeoutError(f"任务轮询超时: task_id={task_id}")

        poll_index += 1
        task_resp = session.get(
            task_url,
            headers=headers,
            timeout=(DEFAULT_POLL_CONNECT_TIMEOUT, DEFAULT_POLL_READ_TIMEOUT),
        )
        task_record = save_http_response(run_dir / f"poll-{poll_index:03d}.json", task_resp)
        task_resp.raise_for_status()
        task = task_record.get("json")
        if not isinstance(task, dict):
            raise RuntimeError(f"任务响应不是 JSON 对象: {task_record}")

        data = task.get("data")
        if not isinstance(data, dict):
            raise RuntimeError(f"任务响应格式异常: {task}")

        status = data.get("status")
        print(f"任务轮询 #{poll_index}: status={status}")
        if status == "completed":
            result = data.get("result", {})
            images = result.get("images", [])
            if not images:
                raise RuntimeError(f"任务已完成，但没有返回图片: {task}")
            return extract_image_bytes_from_item(images[0], session)

        if status == "failed":
            error = data.get("error", {})
            if isinstance(error, dict):
                message = error.get("message") or str(error)
            else:
                message = str(error)
            raise RuntimeError(message or "任务失败")

        time.sleep(poll_interval)


def main() -> None:
    load_env_file()

    api_key = os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY")
    base_url = normalize_base_url(
        os.getenv("OPENAI_BASE_URL") or os.getenv("BASE_URL") or "https://api.toskaxy.xyz"
    )
    if not api_key:
        raise RuntimeError("未找到 API Key，请先在 .env 中设置 OPENAI_API_KEY")

    payload, model, size, output_file, request_mode, reference_source, mask_source = build_payload(base_url)
    resolution = payload.get("resolution")
    run_dir = make_run_dir()
    output_path = resolve_output_path(run_dir, output_file)
    session = build_request_session(base_url)
    request_style = resolve_request_style(base_url)
    request_timeout = env_int("IMAGE_REQUEST_TIMEOUT", DEFAULT_READ_TIMEOUT)
    max_poll_wait = env_int("IMAGE_MAX_POLL_WAIT", DEFAULT_MAX_POLL_WAIT)
    initial_poll_delay = env_int("IMAGE_INITIAL_POLL_DELAY", DEFAULT_INITIAL_POLL_DELAY)
    poll_interval = env_int("IMAGE_POLL_INTERVAL", DEFAULT_POLL_INTERVAL)
    endpoint_path = "/v1/images/edits" if request_mode == "edit" else "/v1/images/generations"
    min_output_edge = resolve_min_output_edge(size, resolution, request_style)

    print(
        "开始生成图片:",
        f"mode={request_mode}",
        f"request_style={request_style}",
        f"model={model}",
        f"size={size}",
        f"resolution={resolution or 'none'}",
        f"request_timeout={request_timeout}s",
        f"max_poll_wait={max_poll_wait}s",
        f"initial_poll_delay={initial_poll_delay}s",
        f"poll_interval={poll_interval}s",
        f"endpoint={base_url}{endpoint_path}",
        f"output={output_path}",
        f"responses={run_dir}",
    )

    dump_json(run_dir / "request.json", {
        "base_url": base_url,
        "endpoint_path": endpoint_path,
        "request_mode": request_mode,
        "model": model,
        "size": size,
        "resolution": resolution,
        "output_file": str(output_path),
        "request_timeout_seconds": request_timeout,
        "max_poll_wait_seconds": max_poll_wait,
        "initial_poll_delay_seconds": initial_poll_delay,
        "poll_interval_seconds": poll_interval,
        "request_style": request_style,
        "min_output_edge": min_output_edge,
        "reference_image_present": bool(reference_source),
        "mask_image_present": bool(mask_source),
        "payload": payload,
    })

    try:
        if request_mode == "edit":
            reference_input = image_source_to_input(reference_source or "", "reference-image", session)
            mask_input = (
                image_source_to_input(mask_source, "mask-image", session)
                if mask_source
                else None
            )
            resp = session.post(
                f"{base_url}{endpoint_path}",
                data=payload_to_form_fields(payload),
                files=build_edit_files(request_style, reference_input, mask_input),
                headers={"Authorization": f"Bearer {api_key}"},
                timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
            )
        else:
            resp = session.post(
                f"{base_url}{endpoint_path}",
                json=payload,
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                timeout=(DEFAULT_CONNECT_TIMEOUT, request_timeout),
            )
        post_record = save_http_response(run_dir / "post-response.json", resp)
        resp.raise_for_status()
        data = post_record.get("json")
        if not isinstance(data, dict):
            raise RuntimeError(f"post response is not a JSON object: {post_record}")

        response_data = data.get("data")
        if isinstance(response_data, list) and response_data:
            image_bytes = extract_image_bytes_from_item(response_data[0], session)
            save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
            return

        if isinstance(response_data, dict) and response_data.get("id"):
            image_bytes = poll_task(
                str(response_data["id"]),
                base_url,
                {"Authorization": f"Bearer {api_key}"},
                session,
                run_dir,
                max_poll_wait,
                initial_poll_delay,
                poll_interval,
            )
            save_image_bytes(image_bytes, output_path, min_output_edge=min_output_edge)
            return

        raise RuntimeError(f"unexpected response: {data}")
    except BaseException as exc:
        save_error_artifact(run_dir / "error.json", exc, {"output_file": str(output_path)})
        raise


if __name__ == "__main__":
    main()