import asyncio
import base64
from contextlib import ExitStack
from datetime import datetime
import importlib
import importlib.util
import json
import mimetypes
import os
import subprocess
import sys
import uuid


def _import_dependencies():
    dependencies = ["tos", "aiohttp"]
    missing = [dep for dep in dependencies if importlib.util.find_spec(dep) is None]

    if not missing:
        return

    print(f"### [Hitem3D] Installing missing dependencies for ComfyUI: {missing}...")

    index_url = "https://pypi.tuna.tsinghua.edu.cn/simple"

    attempts = [
        [sys.executable, "-m", "pip", "install", *missing, "-i", index_url],
        [sys.executable, "-m", "pip", "install", *missing],
        [sys.executable, "-m", "pip", "install", "--user", *missing],
    ]

    success = False
    for cmd in attempts:
        try:
            result = subprocess.run(cmd, check=False)
            if result.returncode == 0:
                success = True
                break
        except Exception:
            pass

    if success:
        importlib.invalidate_caches()
        return

    raise RuntimeError(
        f"Failed to auto-install {missing}. Run manually in a terminal:\n"
        f'"{sys.executable}" -m pip install {" ".join(missing)}'
    )


_import_dependencies()
import aiohttp
import tos
from comfy.model_management import throw_exception_if_processing_interrupted


class Hitem3dAPI:
    API_BASE_URL = "https://api.hitem3d.ai/open-api/v1"
    API_HOST = "api.hitem3d.ai"
    USER_AGENT = "Apifox/1.0.0 (https://apifox.com)"
    REQUEST_TIMEOUT_SECONDS = 30
    TASK_POLL_INTERVAL_SECONDS = 10
    TASK_POLL_TIMEOUT_SECONDS = 60 * 60
    TASK_POLL_REQUEST_TIMEOUT_SECONDS = 30
    TOS_ENDPOINT = "tos-ap-southeast-1.volces.com"
    TOS_REGION = "ap-southeast-1"
    TOS_BUCKET_NAME = "mm-sparc3d-prod"
    TOS_PUBLIC_BASE_URL = "https://hitem3dstatic.zaohaowu.net/"
    VALID_MESH_FORMATS = {".glb", ".usdz", ".obj", ".fbx", ".stl"}
    TASK_PENDING_STATES = {"created", "queueing", "processing"}

    def __init__(self, ak, sk):
        self.ak = ak
        self.sk = sk

    def _build_headers(self, authorization=None, extra_headers=None):
        headers = {
            "User-Agent": self.USER_AGENT,
            "Accept": "*/*",
            "Host": self.API_HOST,
            "Connection": "keep-alive",
        }
        if authorization:
            headers["Authorization"] = authorization
        if extra_headers:
            headers.update(extra_headers)
        return headers

    def _guess_mime_type(self, file_path):
        mime, _ = mimetypes.guess_type(file_path)
        return mime or "image/jpeg"

    def _extract_error_message(self, body, fallback):
        if isinstance(body, dict):
            for key in ("msg", "message", "detail", "error"):
                value = body.get(key)
                if value:
                    return str(value)
        return fallback

    def _build_form_data(self, files, data):
        form_data = aiohttp.FormData()
        for key, value in data.items():
            if value is None:
                continue
            form_data.add_field(key, str(value))

        file_items = files.items() if isinstance(files, dict) else files
        for field_name, file_info in file_items:
            filename, file_obj, content_type = file_info
            form_data.add_field(
                field_name,
                file_obj,
                filename=filename,
                content_type=content_type,
            )
        return form_data

    async def _parse_json_response(self, response, error_message):
        text = await response.text()
        if not text.strip():
            return {}
        try:
            return json.loads(text)
        except json.JSONDecodeError as exc:
            compact_text = " ".join(text.strip().split())[:300]
            raise RuntimeError(
                f"{error_message}: invalid JSON response: {compact_text or '<empty>'}"
            ) from exc

    async def _request_json(
        self,
        session,
        method,
        path,
        authorization=None,
        *,
        params=None,
        data=None,
        timeout_seconds=None,
        extra_headers=None,
        error_message=None,
    ):
        timeout = aiohttp.ClientTimeout(
            total=timeout_seconds or self.REQUEST_TIMEOUT_SECONDS
        )
        url = f"{self.API_BASE_URL}/{path.lstrip('/')}"
        request_error = error_message or f"{method} {path} failed"

        try:
            async with session.request(
                method,
                url,
                headers=self._build_headers(
                    authorization=authorization,
                    extra_headers=extra_headers,
                ),
                params=params,
                data=data,
                timeout=timeout,
            ) as response:
                body = await self._parse_json_response(response, request_error)
                return response.status, body
        except asyncio.TimeoutError as exc:
            raise RuntimeError(f"{request_error}: request timed out") from exc
        except aiohttp.ClientError as exc:
            raise RuntimeError(f"{request_error}: {exc}") from exc

    async def _fetch_access_token(self, session):
        token = f"{self.ak}:{self.sk}".encode("utf-8")
        auth = "Basic " + base64.b64encode(token).decode("ascii")
        status_code, body = await self._request_json(
            session,
            "POST",
            "auth/token",
            authorization=auth,
            extra_headers={"Content-Type": "application/json"},
            data={},
            error_message="get access token failed",
        )
        if status_code != 200 or body.get("code") != 200:
            msg = self._extract_error_message(body, "Unknown error")
            raise RuntimeError(f"get access token failed: {msg}")

        data = body.get("data") or {}
        access_token = data.get("accessToken")
        if not access_token:
            raise RuntimeError("get access token failed: accessToken is missing")
        return access_token

    async def _get_authorization(self, session):
        access_token = await self._fetch_access_token(session)
        return "Bearer " + access_token

    async def _submit_and_wait(self, session, files, data, authorization):
        status_code, body = await self.submit_task(session, files, data, authorization)
        return await self.handle_task_response(
            session,
            status_code,
            body,
            authorization,
        )

    async def get_access_token(self):
        async with aiohttp.ClientSession() as session:
            return await self._fetch_access_token(session)

    async def image_to_3d(self, image_path, request_type, face, model, resolution, pbr=None):
        pbr_value = int(bool(pbr)) if pbr is not None else None
        async with aiohttp.ClientSession() as session:
            authorization = await self._get_authorization(session)
            data = {
                "request_type": request_type,
                "face": face,
                "resolution": resolution,
                "model": model,
                "format": 2,
                "pbr": pbr_value,
            }
            with open(image_path, "rb") as image_file:
                files = {
                    "images": (
                        os.path.basename(image_path),
                        image_file,
                        self._guess_mime_type(image_path),
                    )
                }
                return await self._submit_and_wait(session, files, data, authorization)

    async def multi_view_to_3d(
        self,
        multi_images,
        multi_images_bit,
        request_type,
        face,
        model,
        resolution,
        pbr=None,
    ):
        pbr_value = int(bool(pbr)) if pbr is not None else None
        async with aiohttp.ClientSession() as session:
            authorization = await self._get_authorization(session)
            data = {
                "request_type": request_type,
                "face": face,
                "resolution": resolution,
                "model": model,
                "format": 2,
                "multi_images_bit": multi_images_bit,
                "pbr": pbr_value,
            }
            with ExitStack() as stack:
                files = []
                for img_path in multi_images:
                    if img_path is None:
                        continue
                    filename = os.path.basename(img_path)
                    image_file = stack.enter_context(open(img_path, "rb"))
                    files.append(
                        (
                            "multi_images",
                            (filename, image_file, self._guess_mime_type(img_path)),
                        )
                    )

                if not files:
                    raise ValueError("At least one image is required")

                return await self._submit_and_wait(session, files, data, authorization)

    async def texture(self, image_path, mesh_url, model, pbr=None):
        pbr_value = int(bool(pbr)) if pbr is not None else None
        async with aiohttp.ClientSession() as session:
            authorization = await self._get_authorization(session)
            # Texture node doesn't expose `resolution` in UI. Use model-specific defaults.
            if model == "hitem3dv2.1":
                texture_resolution = "1536fast"
            elif model == "scene-portraitv2.1":
                texture_resolution = "1536profast"
            else:
                texture_resolution = "512"
            data = {
                "request_type": 2,
                "resolution": texture_resolution,
                "model": model,
                "format": 2,
                "mesh_url": mesh_url,
                "pbr": pbr_value,
            }
            with open(image_path, "rb") as image_file:
                files = {
                    "images": (
                        os.path.basename(image_path),
                        image_file,
                        self._guess_mime_type(image_path),
                    )
                }
                return await self._submit_and_wait(session, files, data, authorization)

    async def upload_file(self, mesh_path):
        _, ext = os.path.splitext(os.path.basename(mesh_path))
        ext = ext.lower()
        if ext not in self.VALID_MESH_FORMATS:
            raise ValueError("Invalid mesh file format")

        async with aiohttp.ClientSession() as session:
            authorization = await self._get_authorization(session)
            status_code, body = await self._request_json(
                session,
                "GET",
                "upload/token",
                authorization=authorization,
                error_message="get upload token failed",
            )
            if status_code != 200 or body.get("code") != 200:
                msg = self._extract_error_message(body, "Unknown error")
                raise RuntimeError(f"Upload model failed: {msg}")

        data = body.get("data") or {}
        access_key_id = data.get("accessKeyId")
        secret_access_key = data.get("secretAccessKey")
        session_key = data.get("sessionKey")
        if not access_key_id or not secret_access_key or not session_key:
            raise RuntimeError("Upload model failed: upload credentials are incomplete")

        object_key = f"comfyui/mesh/{datetime.now().strftime('%Y-%m-%d')}/{uuid.uuid4()}{ext}"
        client = tos.TosClientV2(
            access_key_id,
            secret_access_key,
            self.TOS_ENDPOINT,
            self.TOS_REGION,
            security_token=session_key,
        )
        await asyncio.to_thread(
            client.put_object_from_file,
            self.TOS_BUCKET_NAME,
            object_key,
            mesh_path,
        )
        return self.TOS_PUBLIC_BASE_URL + object_key

    async def submit_task(self, session, files, data, authorization):
        payload = dict(data)
        payload["plug_type"] = "COMFYUI"
        payload["plug_task_id"] = uuid.uuid4().hex
        return await self._request_json(
            session,
            "POST",
            "submit-task",
            authorization=authorization,
            data=self._build_form_data(files, payload),
            error_message="submit task failed",
        )

    async def handle_task_response(self, session, status_code, body, authorization):
        if status_code != 200:
            return {
                "status": "error",
                "message": self._extract_error_message(
                    body,
                    f"submit task failed with status code {status_code}",
                ),
                "task_id": None,
            }

        if body.get("code") != 200:
            return {
                "status": "error",
                "message": self._extract_error_message(
                    body,
                    "An unexpected error occurred",
                ),
                "task_id": None,
            }

        task_id = (body.get("data") or {}).get("task_id")
        if not task_id:
            return {
                "status": "error",
                "message": "submit task failed: task_id is missing",
                "task_id": None,
            }

        result = await self.task_status(session, task_id, authorization)
        if isinstance(result, str):
            return {
                "status": "error",
                "message": result,
                "task_id": task_id,
            }

        model_url = result.get("url")
        if not model_url:
            return {
                "status": "error",
                "message": "No downloadable model URL found in task result",
                "task_id": task_id,
            }
        return {"status": "success", "model_url": model_url, "task_id": task_id}

    async def task_status(self, session, task_id, authorization):
        start_time = asyncio.get_running_loop().time()
        last_state = None

        while True:
            throw_exception_if_processing_interrupted()

            if asyncio.get_running_loop().time() - start_time > self.TASK_POLL_TIMEOUT_SECONDS:
                return "Task polling timed out."

            try:
                status_code, body = await self._request_json(
                    session,
                    "GET",
                    "query-task",
                    authorization=authorization,
                    params={"task_id": task_id},
                    timeout_seconds=self.TASK_POLL_REQUEST_TIMEOUT_SECONDS,
                    extra_headers={"Accept-Language": "en-US"},
                    error_message="query task failed",
                )
            except RuntimeError as exc:
                print(f"Polling task retry because of temporary error: {exc}")
                await asyncio.sleep(self.TASK_POLL_INTERVAL_SECONDS)
                continue

            if status_code != 200:
                print(f"Polling task returned unexpected status code: {status_code}")
                await asyncio.sleep(self.TASK_POLL_INTERVAL_SECONDS)
                continue

            if body.get("code") != 200:
                msg = self._extract_error_message(body, "Unknown error")
                return f"Task query failed: {msg}"

            data_obj = body.get("data") or {}
            state = data_obj.get("state")
            if state != last_state:
                print(f"Task {task_id} state: {state}")
                last_state = state

            if state in self.TASK_PENDING_STATES:
                await asyncio.sleep(self.TASK_POLL_INTERVAL_SECONDS)
                continue
            if state == "success":
                return data_obj
            if state == "failed":
                return self._extract_error_message(data_obj, "Task failed.")

            print(f"Polling task returned unexpected state: {state}")
            await asyncio.sleep(self.TASK_POLL_INTERVAL_SECONDS)