#!/usr/bin/env python3
"""Shared authentication, HTTP requests and credit estimation."""
import hashlib
import os
import re
import tempfile
from contextlib import contextmanager
import time
import argparse
import json
from pathlib import Path
import sys
import urllib.error
import urllib.parse
import urllib.request

BASE = "https://openclaw-service.yoooclaw.com/model-proxy/v1"

class ApiError(Exception):
    def __init__(self, message, **details):
        super().__init__(message)
        self.details = details


def error_body(raw, key):
    # Preserve backend diagnostics without echoing this request's credential.
    text = raw.decode("utf-8", errors="replace") if isinstance(raw, bytes) else raw
    text = text.replace(key, "[REDACTED]")
    try:
        return json.loads(text)
    except ValueError:
        return text

class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None

def load_api_key():
    config_path = Path.home() / ".openclaw" / "openclaw.json"
    try:
        config = json.loads(config_path.read_text(encoding="utf-8"))
        key = config["models"]["providers"]["custom_provider_yoooclaw"]["apiKey"]
    except (OSError, ValueError, KeyError, TypeError):
        raise ApiError("无法从 ~/.openclaw/openclaw.json 读取 models.providers.custom_provider_yoooclaw.apiKey，请检查配置。") from None
    if not isinstance(key, str) or not key.strip() or any(c in key for c in "\r\n"):
        raise ApiError("OpenClaw 中的 apiKey 必须是非空字符串且不能包含换行。")
    return key.strip()

def request(method, path, payload=None, *, timeout=None):
    key = load_api_key()
    headers = {"Authorization": "Bearer " + key, "Content-Type": "application/json"}
    body = json.dumps(payload, ensure_ascii=False).encode() if payload is not None else None
    req = urllib.request.Request(BASE + path, data=body, headers=headers, method=method)
    if timeout is None:
        timeout = 300 if method == "POST" and path == "/images/generations" else 60
    try:
        with urllib.request.build_opener(NoRedirect).open(req, timeout=timeout) as response:
            result = json.load(response)
    except urllib.error.HTTPError as exc:
        details = {"http_status": exc.code}
        try:
            raw = exc.read(16385)
            details["response_body"] = error_body(raw[:16384], key)
            if len(raw) > 16384:
                details["response_body_truncated"] = True
        except (OSError, ValueError):
            details["response_body_unavailable"] = True
        finally:
            exc.close()
        if exc.headers and exc.headers.get("x-request-id"):
            details["request_id"] = exc.headers.get("x-request-id").replace(key, "[REDACTED]")
        raise ApiError(f"服务返回 HTTP {exc.code}。未自动重试。", **details) from None
    except (urllib.error.URLError, TimeoutError, OSError):
        raise ApiError("请求未完成，结果不确定。不要重复提交生成请求；已有任务请使用任务 ID 查询。") from None
    except (ValueError, UnicodeError):
        raise ApiError("服务返回无法解析的响应。不要自动重新提交生成请求。") from None
    if not isinstance(result, dict):
        raise ApiError("服务返回的 JSON 结构不符合预期。")
    if result.get("error") or result.get("code") not in (None, 0, "0", 200, "200"):
        raise ApiError("服务返回业务错误；请核实鉴权、参数或余额后再继续。", response_body=error_body(json.dumps(result, ensure_ascii=False), key))
    return result

def positive(value):
    result = int(value)
    if result <= 0:
        raise argparse.ArgumentTypeError("必须是正整数")
    return result

def media_url(item):
    url = item.get("yoooclaw_media_url") if isinstance(item, dict) else None
    if isinstance(url, str) and not any(c.isspace() for c in url):
        try:
            parsed = urllib.parse.urlsplit(url)
            if parsed.scheme in ("http", "https") and parsed.hostname and not parsed.username and not parsed.password:
                return url
        except ValueError:
            pass
    raise ApiError("响应缺少有效的 yoooclaw_media_url；请核实后端结果，不要重新生成。")

def emit(value):
    print(json.dumps(value, ensure_ascii=False), flush=True)

def estimate(payload):
    result = request("GET", "/credits/estimate?" + urllib.parse.urlencode(payload))
    emit({"estimate_response": result})

def execute(run, args):
    try:
        return run(args) or 0
    except (ApiError, OSError, ValueError) as exc:
        message = str(exc) if isinstance(exc, ApiError) else "本地配置读取失败，请检查凭证文件格式和权限。"
        emit({"status": "error", "message": message, **(exc.details if isinstance(exc, ApiError) else {})})
        return 1

def check_generation(args):
    if not args.confirmed:
        raise ApiError("先向用户展示规格与本次积分，获得选择及授权后再传 --confirmed。")
    if not args.prompt.strip():
        raise ApiError("生成描述不能为空。")


def legacy_media_reference(url, kind):
    """Stable opaque delivery reference; plugin renders the actual URL."""
    return "[[" + kind + ":" + hashlib.sha256(url.encode("utf-8")).hexdigest()[:12] + "]]"


MAX_MEDIA_REFERENCES = 128

def media_store_path():
    return Path(os.environ.get("OPENCLAW_STATE_DIR") or (Path.home() / ".openclaw")).expanduser() / "media-references.json"

@contextmanager
def media_store():
    path = media_store_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    lock = path.with_suffix(".lock")
    deadline = time.monotonic() + 10
    while True:
        try:
            lock.mkdir(mode=0o700)
            break
        except FileExistsError:
            if time.monotonic() >= deadline:
                raise ApiError("媒体编号记录正在被占用；请稍后重试原查询或检查残留锁，不要重新生成。")
            time.sleep(0.05)
    temp = None
    try:
        saved = json.loads(path.read_text()) if path.exists() else {}
        if isinstance(saved, dict) and saved.get("version") == 2:
            entries = saved.get("entries")
            counter = saved.get("next_id")
            if type(counter) is not int or counter < 1:
                raise ApiError("媒体编号计数器异常，请核实本地记录。")
        else:
            entries, counter = saved, 1
        if not isinstance(entries, dict) or any(
            not isinstance(row, dict) or not isinstance(row.get("url"), str)
            or row.get("kind") not in ("image", "video")
            or not (re.fullmatch(r"\[\[" + row["kind"] + r":[1-9][0-9]*\]\]", ref)
                    or legacy_media_reference(row["url"], row["kind"]) == ref)
            for ref, row in entries.items()
        ):
            raise ApiError("媒体编号记录格式异常，请核实本地记录，不要重新生成。")
        numeric = [int(ref.split(":")[1][:-2]) for ref in entries
                   if re.fullmatch(r"\[\[(image|video):[1-9][0-9]*\]\]", ref)]
        if numeric and counter <= max(numeric):
            raise ApiError("媒体编号计数器落后于已有记录，请核实本地记录。")
        state = {"version": 2, "next_id": counter, "entries": entries}
        yield state
        while len(entries) > MAX_MEDIA_REFERENCES:
            del entries[next(iter(entries))]
        fd, temp = tempfile.mkstemp(prefix=".media-references-", dir=path.parent)
        with os.fdopen(fd, "w") as output:
            json.dump(state, output, ensure_ascii=False)
            output.flush()
            os.fsync(output.fileno())
        os.replace(temp, path)
    except (OSError, ValueError):
        raise ApiError("媒体编号记录读写失败，请检查状态目录、JSON 格式和权限，不要重新生成。") from None
    finally:
        if temp and os.path.exists(temp):
            os.unlink(temp)
        lock.rmdir()

def remember_media(urls, kind, description=None):
    if kind not in ("image", "video"):
        raise ApiError("媒体类型无效。")
    refs = []
    try:
        with media_store() as state:
            entries = state["entries"]
            for index, url in enumerate(urls, 1):
                ref = next((ref for ref, row in entries.items()
                            if row["url"] == url and row["kind"] == kind
                            and re.fullmatch(r"\[\[" + kind + r":[1-9][0-9]*\]\]", ref)), None)
                if ref is None:
                    ref = "[[" + kind + ":" + str(state["next_id"]) + "]]"
                    state["next_id"] += 1
                previous = entries.pop(ref, {})
                entries[ref] = {"url": url, "kind": kind, "title": description or previous.get("title", previous.get("description", ""))}
                if kind == "image" and len(urls) > 1:
                    entries[ref]["batch_index"] = index
                elif "batch_index" in previous:
                    entries[ref]["batch_index"] = previous["batch_index"]
                refs.append(ref)
    except ApiError as exc:
        exc.details.update(media_urls=urls)
        raise
    return refs

def resolve_image_reference(value):
    if not value.startswith("[["):
        return value
    if not re.fullmatch(r"\[\[image:(?:[1-9][0-9]*|[a-f0-9]{12})\]\]", value):
        raise ApiError("参考图必须使用完整的图片编号 [[image:...]]；视频编号不能作为参考图片。")
    with media_store() as state:
        entries = state["entries"]
        row = entries.get(value)
        if not row or row["kind"] != "image":
            raise ApiError("找不到该图片编号，可能已超出最近 128 个结果或不在当前环境；请提供原图 URL 或本地图片，不要重新生成。")
        entries.pop(value)
        entries[value] = row
        return row["url"]


def media_reference(url, kind):
    """Allocate or reuse a persisted reference under the same cross-process lock."""
    return remember_media([url], kind)[0]


def media_description(value):
    value = value.strip()
    if not value or len(value) > 40 or any(ord(c) < 32 for c in value):
        raise argparse.ArgumentTypeError("请提供 1–40 字符的单行图片或视频标题，用于结果预览和下载链接。")
    return value


def video_description(task_id, description=None):
    # Keep task metadata separate from the bounded media reference cache.
    directory = media_store_path().parent / "media-task-descriptions"
    path = directory / (hashlib.sha256(task_id.encode()).hexdigest() + ".json")
    if description is not None:
        directory.mkdir(parents=True, exist_ok=True, mode=0o700)
        fd, temporary = tempfile.mkstemp(dir=directory)
        try:
            with os.fdopen(fd, "w") as output:
                json.dump({"description": description}, output, ensure_ascii=False)
            os.replace(temporary, path)
        finally:
            if os.path.exists(temporary):
                os.unlink(temporary)
    if not path.exists():
        return None
    return json.loads(path.read_text())["description"]