"""HTTP plumbing shared by every feature module.

Holds the transport concerns only: credentials resolution, headers, retries for
transient failures, automatic JWT refresh from the account's session cookie,
progress polling and file download.  Feature specific request bodies live in
``tripo/features/*``.
"""

import json
import os
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

from . import auth, glb, s3
from .accounts import Accounts
from .config import Config
from .errors import TripoAPIError, TripoTimeout

# Transient transport/server failures worth retrying (the studio occasionally
# returns a bare 502 mid-generation, which must not abort a running poll).
TRANSIENT_STATUS = {429, 500, 502, 503, 504}
RETRY_ATTEMPTS = 4
RETRY_BACKOFF = 1.5

ORIGIN = "https://studio.tripo3d.ai"
USER_AGENT = "tripocli/2.1.0 (Mozilla/5.0; AI Agent Client)"

# Model formats the studio's "上传 3D 模型" importer accepts (UI: OBJ/FBX/STL/GLB ≤150MB).
MODEL_FORMATS = ("glb", "obj", "fbx", "stl")

NO_CREDENTIALS = (
    "未找到可用凭据。请先添加账号（粘贴浏览器 cookie 即可，之后 CLI 自动续期 JWT）：\n"
    "  tripocli account add --cookie \"<整条 Cookie 或 ory_kratos_session 的值>\"\n"
    "也可用环境变量 TRIPO_COOKIE 提供。"
)


def _jwt_expiry(token):
    import base64
    try:
        payload = token.split(".")[1]
        payload += "=" * (-len(payload) % 4)
        return int(json.loads(base64.urlsafe_b64decode(payload).decode()).get("exp") or 0)
    except Exception:
        return 0


class TripoClient:
    def __init__(self, token=None, device_id=None, api_base=None, account=None,
                 auto_refresh=True, accounts=None, config=None):
        self.config = config or Config()
        self.accounts = accounts or Accounts()

        env_token = os.environ.get("TRIPO_TOKEN") or os.environ.get("TRIPO_API_KEY")
        env_cookie = os.environ.get("TRIPO_COOKIE")
        self.account_name, record = self.accounts.get(account or os.environ.get("TRIPO_ACCOUNT"))
        record = record or {}

        self.token = token or env_token or record.get("token")
        self.cookie = env_cookie or record.get("cookie")
        self.device_id = device_id or self.config.device_id
        self.api_base = (api_base or self.config.api_base).rstrip("/")
        self.auto_refresh = auto_refresh and not token and not env_token
        self.token_source = "env" if env_token else ("account" if record.get("token") else None)

    # --- credentials ---

    def token_expiring(self):
        exp = _jwt_expiry(self.token or "")
        return not exp or exp - time.time() < 120

    def refresh_token(self):
        """Mint a fresh JWT from the session cookie (pure HTTP, no browser)."""
        if not self.cookie:
            raise TripoAPIError(NO_CREDENTIALS, code="no-credentials")
        token, session_expires = auth.mint_token(self.api_base, self.cookie, self.device_id)
        self.token = token
        if self.account_name:
            self.accounts.update_token(self.account_name, token, session_expires)
        return True

    def ensure_credentials(self):
        if self.token and not self.token_expiring():
            return
        if self.cookie:
            self.refresh_token()
            return
        if not self.token:
            raise TripoAPIError(NO_CREDENTIALS, code="no-credentials")

    # --- transport ---

    def _headers(self, extra=None):
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Origin": ORIGIN,
            "Referer": ORIGIN + "/",
            "User-Agent": USER_AGENT,
            "x-tripo-device-id": self.device_id,
        }
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        if extra:
            headers.update(extra)
        return headers

    def _raw_request(self, method, endpoint, params=None, body=None, headers=None, timeout=30):
        url = f"{self.api_base}{endpoint}"
        if params:
            query = urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})
            if query:
                url = f"{url}?{query}"

        payload = json.dumps(body).encode("utf-8") if body is not None else None
        request = urllib.request.Request(
            url, data=payload, headers=self._headers(headers), method=method.upper()
        )

        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                text = response.read().decode("utf-8")
                if not text:
                    return None
                return json.loads(text)
        except urllib.error.HTTPError as error:
            err_body = error.read().decode("utf-8", "replace")
            try:
                err_json = json.loads(err_body)
            except json.JSONDecodeError:
                raise TripoAPIError(f"HTTP {error.code}: {error.reason}", code=error.code) from None
            raise TripoAPIError(
                err_json.get("message", f"HTTP {error.code}"),
                code=err_json.get("code", error.code),
                suggestion=err_json.get("suggestion"),
                raw=err_json,
            ) from None
        except urllib.error.URLError as error:
            raise TripoAPIError(f"Network error: {error.reason}") from None

    def request(self, method, endpoint, params=None, body=None, headers=None, timeout=30):
        """Perform an API call, retrying transient failures and expired tokens."""
        attempt = 0
        refreshed = False
        while True:
            attempt += 1
            try:
                result = self._raw_request(method, endpoint, params, body, headers, timeout)
                if isinstance(result, dict) and result.get("code") not in (None, 0):
                    raise TripoAPIError(
                        result.get("message", "API request failed"),
                        code=result.get("code"),
                        suggestion=result.get("suggestion"),
                        raw=result,
                    )
                return result
            except TripoAPIError as error:
                if error.is_auth_error and self.auto_refresh and not refreshed:
                    refreshed = True
                    try:
                        if self.refresh_token():
                            continue
                    except TripoAPIError:
                        pass  # fall through and surface the original error
                transient = error.code in TRANSIENT_STATUS or error.code is None
                if transient and attempt < RETRY_ATTEMPTS:
                    time.sleep(RETRY_BACKOFF * attempt)
                    continue
                raise

    # --- shared API calls ---

    def get_data(self, method, endpoint, **kwargs):
        result = self.request(method, endpoint, **kwargs)
        return result.get("data") if result else None

    def progress(self, operator_ids):
        if isinstance(operator_ids, str):
            operator_ids = [operator_ids]
        return self.get_data("POST", "/v2/studio/progress", body={"ids": operator_ids}) or []

    def project_detail(self, project_id, locale="zh-CN"):
        return self.get_data("GET", f"/v2/studio/project/detail/v3/{project_id}", params={"locale": locale})

    def project_operator(self, project_id):
        """Current operator block of a project (carries operator_id + model_version)."""
        detail = self.project_detail(project_id) or {}
        return detail.get("operator") or {}

    def part_names(self, project_id):
        """Node names inside the project's current mesh (the ``part_name_list``).

        Read from the model file itself, since the detail API does not expose
        them.  Returns an empty list when the model cannot be read.
        """
        detail = self.project_detail(project_id) or {}
        url = (detail.get("operator") or {}).get("model_url")
        if not url:
            return []
        request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(request, timeout=120) as response:
            data = response.read()
        return glb.node_names(data)

    def payment_profile(self):
        return self.get_data("GET", "/v2/studio/user/profile/payment")

    def whoami(self):
        return self.get_data("GET", "/v2/studio/studio/whoami", params={"tokenizeAs": "default_jwt"})

    def quota(self):
        return self.get_data("POST", "/v2/studio/operation/quota", body={})

    def assets(self, asset_type="mine", type="model", page=1, size=20, offset=0, locale="zh-CN"):
        return self.get_data(
            "GET",
            "/v2/studio/assets/v2",
            params={
                "asset_type": asset_type,
                "type": type,
                "page": page,
                "size": size,
                "offset": offset,
                "locale": locale,
            },
        )

    def download_url(self, operator_id, file_name):
        data = self.get_data(
            "POST", "/v2/studio/operation/download_with_name",
            body={"operator_id": operator_id, "file_name": file_name},
        )
        return (data or {}).get("model_url")

    def export(self, project_id, operator_id, export_format, name="model", model_version="default"):
        fmt = {"glb": "gltf", "gltf": "gltf"}.get(export_format.lower(), export_format.lower())
        return self.get_data(
            "POST", "/v2/studio/operation/export",
            body={
                "project_id": project_id,
                "operator_id": operator_id,
                "format": fmt,
                "name": name,
                "model_version": model_version,
            },
        )

    # --- images ---

    def upload_image(self, path):
        """Upload a local image and return an image object ready to be submitted.

        Verified flow: temporary_token -> SigV4 PUT -> ``/audit/image`` -> the
        ``image`` object must carry ``image_audit_result`` and ``image_source``,
        otherwise every generation endpoint answers 11002.
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"图片文件不存在: {path}")

        fmt = path.suffix.lower().lstrip(".") or "png"
        if fmt == "jpeg":
            fmt = "jpg"
        if fmt not in ("png", "jpg", "webp"):
            raise ValueError(f"不支持的图片格式: .{fmt} (支持 png/jpg/webp)")
        content_type = "image/jpeg" if fmt == "jpg" else f"image/{fmt}"

        sts = self.get_data(
            "POST", "/v2/studio/storage/temporary_token", body={"client": "aws", "format": fmt}
        )
        s3.upload_bytes(sts, path.read_bytes(), content_type=content_type)

        image = {"bucket": sts["resource_bucket"], "key": sts["resource_uri"]}
        audit = self.get_data("POST", "/v2/studio/audit/image", body={"image": image}) or {}
        image["image_audit_result"] = audit.get("result", "pass")
        image["image_source"] = "upload"
        return image

    # --- models ---

    def upload_model(self, path):
        """Upload a local 3D model and return the ``import_user_model`` payload.

        Verified flow: ``temporary_token`` (format = file extension) -> SigV4 PUT
        -> the import body is ``{format, model:{bucket,key}, name,
        transform_matrix, use_original_uv}``.  The studio UI rotates the mesh
        before importing (that is what ``transform_matrix`` carries); a CLI
        import uses the identity matrix.
        """
        path = Path(path)
        if not path.exists():
            raise FileNotFoundError(f"模型文件不存在: {path}")

        fmt = path.suffix.lower().lstrip(".")
        if fmt not in MODEL_FORMATS:
            raise ValueError(f"不支持的模型格式: .{fmt} (支持 glb/obj/fbx/stl)")

        sts = self.get_data(
            "POST", "/v2/studio/storage/temporary_token", body={"client": "aws", "format": fmt}
        )
        s3.upload_bytes(sts, path.read_bytes(), content_type="application/octet-stream")
        return {
            "format": fmt,
            "model": {"bucket": sts["resource_bucket"], "key": sts["resource_uri"]},
            "name": path.name,
        }

    # --- waiting & downloading ---

    def wait_task(self, operator_id, timeout=900, poll_interval=3, on_progress=None):
        """Poll until the operator succeeds, fails or the timeout expires."""
        if not operator_id:
            # Guard: polling a missing id would otherwise spin until the timeout.
            raise TripoAPIError("没有可等待的 operator_id（接口未返回任务 id）", code="no-operator")
        deadline = time.time() + timeout
        last_progress = -1
        while time.time() < deadline:
            items = self.progress([operator_id])
            if items:
                task = items[0]
                status = task.get("status")
                progress = task.get("progress", 0)
                if progress != last_progress:
                    last_progress = progress
                    if on_progress:
                        on_progress(status, progress, task.get("left_time", -1), task)
                if status == "success":
                    return task
                if status in ("failed", "banned", "expired"):
                    raise TripoAPIError(f"任务 {operator_id} 失败: {json.dumps(task, ensure_ascii=False)}")
            time.sleep(poll_interval)
        raise TripoTimeout(f"任务 {operator_id} 在 {timeout}s 内未完成")

    def download(self, url, output_path):
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        request = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(request, timeout=300) as response, open(output_path, "wb") as handle:
            while True:
                chunk = response.read(64 * 1024)
                if not chunk:
                    break
                handle.write(chunk)
        return output_path