"""Generate and download a base mesh from the Hunyuan 3D website."""
from __future__ import annotations

import json
import re
import time
import zipfile
from pathlib import Path
from typing import Callable, Optional

from .auth import HUNYUAN_URL, is_logged_in, storage_state_path

FACE_LABEL = "50k"
FACE_LABELS = ("50k", "500k", "1m")
GENERATION_TIMEOUT_SEC = 1200
POLL_INTERVAL_SEC = 5
API_LIST = "https://3d.hunyuan.tencent.com/api/3d/creations/list"
API_DETAIL = "https://3d.hunyuan.tencent.com/api/3d/creations/detail"

UpdateFn = Callable[..., None]


def _browser_channel() -> str:
    import os

    return os.environ.get("HUNYUAN_BROWSER_CHANNEL", "msedge").strip() or "msedge"


def _click_visible_text(page, labels, exact=False, timeout=5000):
    for label in labels:
        locator = page.get_by_text(label, exact=exact)
        try:
            total = locator.count()
        except Exception:
            total = 0
        for index in range(total):
            target = locator.nth(index)
            try:
                if not target.is_visible():
                    continue
                target.scroll_into_view_if_needed(timeout=2000)
                target.click(timeout=timeout)
                page.wait_for_timeout(400)
                return True
            except Exception:
                continue
    return False


def _dismiss_popups(page, log=None):
    """Close common confirm/agreement toasts that block generate submit."""
    for label in ("我知道了", "知道了", "同意并继续", "同意", "确认", "继续", "关闭"):
        if _click_visible_text(page, (label,), exact=False, timeout=1200):
            if log:
                log.write(f"dismissed popup via 「{label}」\n")
            page.wait_for_timeout(400)


def _click_generate(page, log=None):
    """Click the primary studio generate button (not nav text)."""
    _dismiss_popups(page, log)
    locators = (
        page.get_by_role("button", name=re.compile(r"立即生成")),
        page.locator('button.t-button--theme-primary:has-text("立即生成")'),
        page.locator('button:has-text("立即生成")'),
        page.get_by_text("立即生成", exact=False),
    )
    for locator in locators:
        try:
            total = locator.count()
        except Exception:
            continue
        for index in range(total):
            target = locator.nth(index)
            try:
                if not target.is_visible():
                    continue
                try:
                    if target.is_disabled():
                        continue
                except Exception:
                    pass
                box = target.bounding_box()
                if box and (box["width"] < 80 or box["height"] < 28):
                    continue
                target.scroll_into_view_if_needed(timeout=2000)
                target.click(timeout=5000)
                page.wait_for_timeout(600)
                if log:
                    log.write("clicked generate\n")
                return True
            except Exception as exc:
                if log:
                    log.write(f"generate click skipped: {exc}\n")
                continue
    for index in range(page.get_by_text("立即生成", exact=False).count()):
        target = page.get_by_text("立即生成", exact=False).nth(index)
        try:
            if not target.is_visible():
                continue
            target.click(timeout=4000, force=True)
            if log:
                log.write("clicked generate (force)\n")
            return True
        except Exception:
            continue
    return False


def _wait_generations_api(generations_bodies, page, log, get_creations_id=None, timeout_sec=25):
    deadline = time.monotonic() + timeout_sec
    while time.monotonic() < deadline:
        if generations_bodies or (get_creations_id and get_creations_id()):
            return True
        _dismiss_popups(page, log)
        page.wait_for_timeout(800)
    return bool(generations_bodies or (get_creations_id and get_creations_id()))


def _in_studio(page) -> bool:
    try:
        return bool(
            page.get_by_text("模型面数", exact=True).count()
            and page.get_by_text("立即生成", exact=False).count()
        )
    except Exception:
        return False


def _enter_studio(page, log):
    if _in_studio(page):
        log.write("already in studio\n")
        return
    if _click_visible_text(page, ("立即开始",), exact=False):
        page.wait_for_timeout(2500)
    if not _in_studio(page):
        _click_visible_text(page, ("图/文生3D",), exact=False)
        page.wait_for_timeout(800)
        _click_visible_text(page, ("立即开始", "开始"), exact=False)
        page.wait_for_timeout(2500)
    deadline = time.monotonic() + 30
    while time.monotonic() < deadline and not _in_studio(page):
        page.wait_for_timeout(500)
    if not _in_studio(page):
        raise RuntimeError("未能进入混元「图/文生3D」创作页，请确认登录态后重试。")
    log.write("entered studio\n")


def _select_face_count(page, label=FACE_LABEL, log=None):
    aliases = {
        "50k": ("50k", "50K"),
        "500k": ("500k", "500K"),
        "1m": ("1m", "1M"),
    }
    candidates = aliases.get(label, (label,))
    for candidate in candidates:
        pattern = re.compile(rf"^\s*{re.escape(candidate)}\s*$", re.I)
        locator = page.get_by_text(pattern)
        try:
            total = locator.count()
        except Exception:
            total = 0
        for index in range(total):
            target = locator.nth(index)
            try:
                if not target.is_visible():
                    continue
                box = target.bounding_box()
                if not box or box["width"] > 120 or box["height"] > 80:
                    continue
                target.scroll_into_view_if_needed(timeout=2000)
                target.click(timeout=4000)
                page.wait_for_timeout(400)
                if log:
                    log.write(f"selected face count {candidate}\n")
                return True
            except Exception as exc:
                if log:
                    log.write(f"face click skipped: {exc}\n")
                continue
    if log:
        log.write(f"visible face chip {label} not clicked\n")
    return False


def _dump_debug(page, directory: Path, log, network_hits):
    try:
        page.screenshot(path=str(directory / "hunyuan_debug.png"), full_page=True)
        log.write("saved hunyuan_debug.png\n")
    except Exception as exc:
        log.write(f"screenshot failed: {exc}\n")
    try:
        (directory / "hunyuan_debug.html").write_text(page.content(), encoding="utf-8")
        log.write("saved hunyuan_debug.html\n")
    except Exception as exc:
        log.write(f"html dump failed: {exc}\n")
    (directory / "hunyuan_network.json").write_text(
        json.dumps(network_hits[-120:], ensure_ascii=False, indent=2), encoding="utf-8"
    )


def _api_json(request_context, method, url, **kwargs):
    response = getattr(request_context, method)(url, **kwargs)
    if response.status >= 400:
        if response.status in (401, 403):
            raise PermissionError(f"混元接口鉴权失败 {response.status}: {url}")
        raise RuntimeError(f"混元接口失败 {response.status}: {url}")
    try:
        return response.json()
    except Exception as exc:
        raise RuntimeError(f"混元接口返回非 JSON：{url}") from exc


def _list_creations(request_context):
    return _api_json(
        request_context,
        "post",
        API_LIST,
        data=json.dumps({}),
        headers={"content-type": "application/json"},
    )


def _detail_creation(request_context, creations_id):
    return _api_json(
        request_context, "get", f"{API_DETAIL}?creationsId={creations_id}"
    )


def _creation_progress(creation: dict, page=None):
    """Return (stage_text, progress_0_100, done, failed, glb_url)."""
    status = str(creation.get("status") or "").lower()
    result = (creation.get("result") or [{}])[0] if creation.get("result") else {}
    result_status = str(result.get("status") or status).lower()
    urls = result.get("urlResult") or {}
    glb = urls.get("glb") or urls.get("textureGlb")
    intermediate = result.get("intermediate_outputs") or {}
    has_geometry = bool(
        (intermediate.get("geometry") or {}).get("glb_url") or urls.get("geometryGlb")
    )

    ui_percent = None
    ui_stage = None
    if page is not None:
        try:
            body = page.inner_text("body")
            match = re.search(
                r"(几何生成中|纹理生成中|PBR生成中|生成中)[^\d%]{0,8}(\d{1,3})\s*%", body
            )
            if match:
                ui_stage, ui_percent = match.group(1), int(match.group(2))
        except Exception:
            pass

    if result_status in {"fail", "failed", "error"} or status in {
        "fail",
        "failed",
        "error",
    }:
        return "混元生成失败", 100, False, True, None
    if result_status in {"success", "done", "finished"} and glb:
        return "混元生成完成", 95, True, False, glb

    if ui_percent is not None:
        stage = f"资产页 · {ui_stage} {ui_percent}%"
        if ui_stage and ("纹理" in ui_stage or "PBR" in ui_stage):
            progress = 80 + min(14, ui_percent * 14 // 100)
        else:
            progress = 55 + min(24, ui_percent * 24 // 100)
        return stage, progress, False, False, None

    if has_geometry:
        return "资产页 · 纹理/PBR 生成中", 82, False, False, None
    if status in {"wait", "waiting", "queue", "queued"}:
        return "资产页 · 排队中", 58, False, False, None
    return "资产页 · 几何生成中", 65, False, False, None


def _extract_creations_id(payload):
    if payload is None:
        return None
    if isinstance(payload, dict):
        for key in ("creationsId", "creationId", "CreationsId", "id"):
            value = payload.get(key)
            if isinstance(value, str) and len(value) >= 8:
                return value
        for key in ("data", "result", "creation", "creations"):
            nested = payload.get(key)
            found = _extract_creations_id(nested)
            if found:
                return found
        if isinstance(payload.get("list"), list):
            found = _extract_creations_id(payload.get("list"))
            if found:
                return found
    if isinstance(payload, list):
        for item in payload:
            found = _extract_creations_id(item)
            if found:
                return found
    return None


def _find_new_creation(listing, before_ids):
    creations = listing.get("creations") or []
    for item in creations:
        item_id = item.get("id")
        if item_id and item_id not in before_ids:
            return item
    return None


def _wait_for_new_creations_id(
    request_context, before_ids, captured_id, update, log, timeout_sec=120
):
    if captured_id and captured_id not in before_ids:
        log.write(f"creationsId from generations response: {captured_id}\n")
        return captured_id
    if captured_id and captured_id in before_ids:
        log.write(f"ignored stale creationsId from response: {captured_id}\n")

    deadline = time.monotonic() + timeout_sec
    while time.monotonic() < deadline:
        waited = int(timeout_sec - (deadline - time.monotonic()))
        update(
            event="stage",
            stage=f"等待资产页出现新任务（{waited}s / {timeout_sec}s）",
            progress=54,
        )
        listing = _list_creations(request_context)
        newest = _find_new_creation(listing, before_ids)
        if newest and newest.get("id"):
            log.write(
                f"creationsId from list delta: {newest['id']} status={newest.get('status')} "
                f"total={listing.get('totalCount')}\n"
            )
            return newest["id"]
        time.sleep(POLL_INTERVAL_SEC)
    raise RuntimeError(
        "提交后资产列表未出现新任务。可能未真正提交成功，请查看官网资产页后重试。"
    )


def _download_url(request_context, url, destination: Path, log):
    response = request_context.get(url)
    if response.status >= 400:
        raise RuntimeError(f"下载模型失败 HTTP {response.status}")
    destination.write_bytes(response.body())
    log.write(f"downloaded bytes={destination.stat().st_size} -> {destination}\n")
    return destination


def _extract_mesh(archive_or_file: Path, directory: Path) -> Path:
    if archive_or_file.suffix.lower() == ".glb":
        destination = directory / "generated.glb"
        if archive_or_file.resolve() != destination.resolve():
            destination.write_bytes(archive_or_file.read_bytes())
        return destination
    if archive_or_file.suffix.lower() == ".zip" or zipfile.is_zipfile(archive_or_file):
        extract_to = directory / "hunyuan_extract"
        extract_to.mkdir(exist_ok=True)
        with zipfile.ZipFile(archive_or_file, "r") as archive:
            archive.extractall(extract_to)
        glbs = sorted(extract_to.rglob("*.glb"))
        if not glbs:
            raise RuntimeError("混元下载包中未找到 GLB 模型。")
        destination = directory / "generated.glb"
        destination.write_bytes(glbs[0].read_bytes())
        return destination
    raise RuntimeError(f"不支持的混元下载格式：{archive_or_file.suffix}")


def _poll_until_ready(
    request_context, page, creations_id, update, log, network_hits, directory
):
    deadline = time.monotonic() + GENERATION_TIMEOUT_SEC
    last_stage = ""
    while time.monotonic() < deadline:
        waited = int(GENERATION_TIMEOUT_SEC - (deadline - time.monotonic()))
        try:
            detail = _detail_creation(request_context, creations_id)
        except PermissionError:
            raise
        except Exception:
            listing = _list_creations(request_context)
            detail = next(
                (
                    c
                    for c in (listing.get("creations") or [])
                    if c.get("id") == creations_id
                ),
                None,
            )
            if detail is None:
                detail = {"status": "waiting", "result": []}
        (directory / "hunyuan_creation.json").write_text(
            json.dumps(detail, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        stage, progress, done, failed, glb = _creation_progress(detail, page=page)
        message = f"{stage}（已等待 {waited}s / 最长 {GENERATION_TIMEOUT_SEC}s）"
        if message != last_stage:
            log.write(message + "\n")
            last_stage = message
        update(event="stage", stage=message, progress=progress)
        if failed:
            raise RuntimeError("混元资产页显示生成失败。")
        if done and glb:
            return glb
        try:
            if page.get_by_text("几何生成中", exact=False).count() == 0:
                _click_visible_text(page, ("资产",), exact=True)
        except Exception:
            pass
        page.wait_for_timeout(POLL_INTERVAL_SEC * 1000)
    _dump_debug(page, directory, log, network_hits)
    raise TimeoutError(f"混元资产生成超时（已等待 {GENERATION_TIMEOUT_SEC}s）。")


def _page_looks_logged_out(page) -> bool:
    try:
        if page.get_by_text("登录", exact=True).count():
            login = page.get_by_text("登录", exact=True).first
            if login.is_visible(timeout=500):
                return True
    except Exception:
        pass
    return False


def generate_from_image(
    image_path: Path,
    directory: Path,
    update: UpdateFn,
    face_label: str = FACE_LABEL,
) -> Path:
    """图生3D，返回 directory/generated.glb。鉴权失败抛 PermissionError。"""
    if not is_logged_in():
        raise PermissionError("尚未微信登录混元3D，请先登录。")
    if face_label not in FACE_LABELS:
        raise ValueError(f"面数须为 {FACE_LABELS} 之一，收到: {face_label}")

    from playwright.sync_api import sync_playwright

    image_path = Path(image_path).expanduser().resolve()
    if not image_path.is_file():
        raise FileNotFoundError(f"参考图不存在: {image_path}")

    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    log_path = directory / "generation.log"
    network_hits = []
    api_bodies = {}
    generations_bodies = []
    creations_id = None
    state_path = storage_state_path()

    def on_response(response):
        nonlocal creations_id
        url = response.url
        if "/api/3d/" not in url:
            return
        entry = {"url": url, "status": response.status}
        network_hits.append(entry)
        path = url.split("?", 1)[0]
        if any(
            token in path
            for token in (
                "/creations/generations",
                "/creations/detail",
                "/creations/list",
                "/resource/download",
            )
        ):
            try:
                body = response.json()
                entry["body_keys"] = (
                    list(body)[:20] if isinstance(body, dict) else type(body).__name__
                )
                api_bodies.setdefault(path, body)
                if "/creations/generations" in path and isinstance(body, dict):
                    generations_bodies.append(body)
                    found = _extract_creations_id(body)
                    if found:
                        creations_id = found
                        entry["creationsId"] = found
            except Exception:
                pass

    update(event="stage", stage="打开混元3D 网页", progress=20)
    with log_path.open("w", encoding="utf-8") as log, sync_playwright() as playwright:
        browser = playwright.chromium.launch(
            channel=_browser_channel(), headless=True
        )
        context = browser.new_context(
            storage_state=str(state_path),
            locale="zh-CN",
            accept_downloads=True,
            viewport={"width": 1600, "height": 1200},
        )
        page = context.new_page()
        page.on("response", on_response)
        try:
            page.goto(HUNYUAN_URL, wait_until="domcontentloaded", timeout=60000)
            page.wait_for_timeout(2500)
            log.write("opened hunyuan home\n")

            if _page_looks_logged_out(page):
                raise PermissionError("混元登录凭证已失效，需要重新微信登录。")

            update(event="stage", stage="进入图/文生3D", progress=26)
            _enter_studio(page, log)

            before = _list_creations(context.request)
            before_ids = {
                item.get("id")
                for item in (before.get("creations") or [])
                if item.get("id")
            }
            before_count = before.get("totalCount")
            log.write(
                f"before_submit totalCount={before_count} known_ids={len(before_ids)}\n"
            )

            update(event="stage", stage="切换到图生3D", progress=30)
            _click_visible_text(page, ("图生3D",), exact=True)
            _click_visible_text(page, ("单张图片",), exact=True)

            update(event="stage", stage="上传参考图片", progress=36)
            file_input = page.locator('input[type="file"]').first
            if not file_input.count():
                _click_visible_text(page, ("上传图片",), exact=False)
                file_input = page.locator('input[type="file"]').first
            if not file_input.count():
                _dump_debug(page, directory, log, network_hits)
                raise RuntimeError("未找到图片上传控件，页面结构可能已变化。")
            file_input.set_input_files(str(image_path))
            page.wait_for_timeout(3500)
            log.write(f"uploaded {image_path}\n")
            upload_deadline = time.monotonic() + 30
            while time.monotonic() < upload_deadline:
                if any("/resource/review" in (e.get("url") or "") for e in network_hits):
                    break
                page.wait_for_timeout(500)
            page.wait_for_timeout(1500)

            update(event="stage", stage=f"选择 {face_label} 原始面数", progress=44)
            if not _select_face_count(page, face_label, log=log):
                _dump_debug(page, directory, log, network_hits)
                raise RuntimeError(f"未能点选模型面数 {face_label}。")

            update(event="stage", stage="提交混元生成", progress=52)
            submitted = False
            for attempt in range(1, 4):
                if not _click_generate(page, log=log):
                    _dump_debug(page, directory, log, network_hits)
                    raise RuntimeError("未找到「立即生成」按钮，页面结构可能已变化。")
                log.write(f"generate attempt {attempt}/3\n")
                if _wait_generations_api(
                    generations_bodies,
                    page,
                    log,
                    get_creations_id=lambda: creations_id,
                    timeout_sec=20,
                ):
                    submitted = True
                    log.write("generations API observed\n")
                    break
                log.write("no /creations/generations yet; will retry click\n")
                page.wait_for_timeout(1200)
            if not submitted:
                _dump_debug(page, directory, log, network_hits)
                raise RuntimeError(
                    "点击「立即生成」后官网未发起生成请求（无 /creations/generations）。"
                    "请打开混元官网确认额度/协议弹窗后重试。"
                )

            (directory / "hunyuan_generations.json").write_text(
                json.dumps(generations_bodies, ensure_ascii=False, indent=2)[:200000],
                encoding="utf-8",
            )
            (directory / "hunyuan_api_bodies.json").write_text(
                json.dumps(
                    {
                        k: (
                            v
                            if not str(k).rstrip("/").endswith("list")
                            else {
                                "totalCount": v.get("totalCount"),
                                "first_id": (v.get("creations") or [{}])[0].get("id"),
                            }
                        )
                        for k, v in api_bodies.items()
                    },
                    ensure_ascii=False,
                    indent=2,
                )[:200000],
                encoding="utf-8",
            )

            creations_id = _wait_for_new_creations_id(
                context.request, before_ids, creations_id, update, log, timeout_sec=120
            )
            if creations_id in before_ids:
                raise RuntimeError(f"拒绝绑定旧资产 {creations_id}，请重试提交。")
            log.write(f"bound creationsId={creations_id}\n")
            (directory / "hunyuan_binding.json").write_text(
                json.dumps(
                    {
                        "creationsId": creations_id,
                        "before_count": before_count,
                        "before_ids_sample": sorted(before_ids)[:5],
                    },
                    ensure_ascii=False,
                    indent=2,
                ),
                encoding="utf-8",
            )

            update(event="stage", stage="打开资产页并等待完成", progress=55)
            _click_visible_text(page, ("资产",), exact=True)
            page.wait_for_timeout(1500)

            glb_url = _poll_until_ready(
                context.request, page, creations_id, update, log, network_hits, directory
            )

            update(event="stage", stage="下载混元 GLB", progress=96)
            saved = directory / "hunyuan_raw.glb"
            _download_url(context.request, glb_url, saved, log)
            mesh = _extract_mesh(saved, directory)
            log.write(f"extracted {mesh}\n")
            (directory / "hunyuan_network.json").write_text(
                json.dumps(network_hits[-120:], ensure_ascii=False, indent=2),
                encoding="utf-8",
            )
            context.storage_state(path=str(state_path))
            return mesh
        except Exception as exc:
            log.write(f"ERROR: {exc}\n")
            try:
                _dump_debug(page, directory, log, network_hits)
            except Exception:
                pass
            raise
        finally:
            browser.close()