#!/usr/bin/env python3
"""Python 3.10+, standard library. Default: offline preview. MIT license."""
import argparse
import ast
import base64
import datetime
import hashlib
import ipaddress
import json
import math
import os
from pathlib import Path
import re
import struct
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import HTTPRedirectHandler, Request, build_opener
import uuid

API = "https://api.quickrouter.ai/v1"
FIXTURE = "def calculate_total(subtotal, shipping, discount):\n    return subtotal + shipping + discount\n"
CASES = [(42.5, 7.5, 10, 40), (10, 0, 2, 8), (20, 0, 0, 20), (0, 0, 0, 0)]
EDIT_PROMPT = "Keep the subject's color, shape, details, quantity and proportions unchanged. Replace only the background with a clean cream studio wall and soft side light. Do not add text or new objects."
CODE_PROMPT = '修复下面的 Python 函数，实付金额应等于商品小计 + 运费 - 优惠。保持函数名与三个参数不变，只使用一个 return 和参数的加减表达式。不要导入模块、调用函数或添加类型注解。只返回 JSON，格式 {"code":"完整函数源码"}，不要 Markdown。失败用例 calculate_total(42.5, 7.5, 10) 预期 40，目前返回 60。源码：\n' + FIXTURE


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


def safe_url(value):
    if not isinstance(value, str):
        raise ValueError("Missing HTTPS download URL")
    parsed = urlparse(value)
    if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password:
        raise ValueError("Download requires HTTPS without URL credentials")
    if parsed.hostname.lower() in {"localhost", "localhost.localdomain"} or parsed.hostname.lower().endswith(".local"):
        raise ValueError("Local download hosts are not allowed")
    try:
        address = ipaddress.ip_address(parsed.hostname)
    except ValueError:
        address = None
    if address and not address.is_global:
        raise ValueError("Private download addresses are not allowed")
    return value


def redact(value, key):
    if isinstance(value, str):
        text = value.replace(key, "[REDACTED]") if key else value
        return re.sub(r"sk-[A-Za-z0-9_-]+", "[REDACTED]", text)
    if isinstance(value, dict):
        return {name: redact(item, key) for name, item in value.items()}
    if isinstance(value, list):
        return [redact(item, key) for item in value]
    return value


def read_response(request, limit, timeout=180):
    try:
        with build_opener(NoRedirect()).open(request, timeout=timeout) as response:
            raw = response.read(limit + 1)
            if len(raw) > limit:
                raise ValueError("Response exceeds the example size limit")
            return raw
    except HTTPError as error:
        raise ValueError(f"HTTP {error.code}; inspect console and endpoint. No automatic resubmission.") from None
    except (URLError, TimeoutError):
        raise ValueError("Request failed or timed out. Check console before retrying; the request may have succeeded.") from None


def request_json(url, body, key, content_type="application/json", read=read_response):
    request = Request(url, data=body, method="POST", headers={
        "Authorization": "Bearer " + key, "Content-Type": content_type, "Accept": "application/json",
    })
    parsed = json.loads(read(request, 32 * 1024 * 1024))
    if not isinstance(parsed, dict) or parsed.get("error"):
        raise ValueError("API did not return a successful JSON object; inspect console")
    return parsed


def png_dimensions(data):
    if len(data) < 33 or data[:8] != b"\x89PNG\r\n\x1a\n" or data[12:16] != b"IHDR":
        raise ValueError("Input must be a PNG file, not an extension-renamed file")
    width, height = struct.unpack(">II", data[16:24])
    if not width or not height or width * height > 80_000_000:
        raise ValueError("PNG dimensions are invalid or too large for this example")
    return width, height


def image_extension(data):
    if data.startswith(b"\x89PNG\r\n\x1a\n"):
        png_dimensions(data)
        return "png"
    if len(data) > 4 and data.startswith(b"\xff\xd8\xff") and data.endswith(b"\xff\xd9"):
        return "jpg"
    if len(data) > 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "webp"
    raise ValueError("Result is not a recognized PNG, JPEG or WebP image; no successful image recorded")


def multipart(image, model, boundary):
    png_dimensions(image)
    fields = {"model": model, "prompt": EDIT_PROMPT, "n": "1", "size": "1024x1024", "quality": "low"}
    parts = []
    for name, value in fields.items():
        parts.append(f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n{value}\r\n'.encode())
    parts.append(f'--{boundary}\r\nContent-Disposition: form-data; name="image"; filename="input.png"\r\nContent-Type: image/png\r\n\r\n'.encode() + image + b"\r\n")
    parts.append(f"--{boundary}--\r\n".encode())
    return b"".join(parts), "multipart/form-data; boundary=" + boundary


def code_payload(model):
    return {"model": model, "messages": [{"role": "user", "content": CODE_PROMPT}], "max_tokens": 4096, "stream": False}


def code_acceptance(source):
    if not isinstance(source, str) or len(source) > 4000:
        raise ValueError("Expected a short function source string")
    tree = ast.parse(source)
    if len(tree.body) != 1 or not isinstance(tree.body[0], ast.FunctionDef):
        raise ValueError("Candidate must contain exactly one function")
    fn = tree.body[0]
    args = fn.args
    names = [item.arg for item in args.args]
    if fn.name != "calculate_total" or names != ["subtotal", "shipping", "discount"] or args.posonlyargs or args.kwonlyargs or args.vararg or args.kwarg or args.defaults or fn.decorator_list or fn.returns or any(item.annotation for item in args.args):
        raise ValueError("Candidate changed the required function signature")
    if len(fn.body) != 1 or not isinstance(fn.body[0], ast.Return):
        raise ValueError("Candidate must use one return expression")

    def calculate(node, values):
        if isinstance(node, ast.Name) and node.id in values:
            return values[node.id]
        if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Sub)):
            left, right = calculate(node.left, values), calculate(node.right, values)
            return left + right if isinstance(node.op, ast.Add) else left - right
        raise ValueError("Only parameter names and addition/subtraction are accepted; arbitrary code is never executed")

    results = []
    for subtotal, shipping, discount, expected in CASES:
        actual = calculate(fn.body[0].value, dict(zip(names, (subtotal, shipping, discount))))
        results.append({"input": [subtotal, shipping, discount], "expected": expected, "actual": actual, "passed": math.isclose(actual, expected, abs_tol=1e-9)})
    return {"passed": all(item["passed"] for item in results), "tests": results, "scope": "Four fixed arithmetic cases; no repository tools or arbitrary code execution"}


def save_json(path, value):
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")


def prepare_out(path):
    path.mkdir(parents=True, exist_ok=False)
    save_json(path / "record.json", {"status": "submission_unconfirmed", "billed_usd": None, "observed_at": datetime.datetime.now(datetime.timezone.utc).isoformat()})


def execute_code(model, key, out, request=request_json):
    prepare_out(out)
    (out / "fixture-before.py").write_text(FIXTURE, encoding="utf-8")
    payload = code_payload(model)
    save_json(out / "request.json", payload)
    started = time.perf_counter()
    response = request(API + "/chat/completions", json.dumps(payload).encode(), key)
    save_json(out / "response.json", redact(response, key))
    choices = response.get("choices")
    if not isinstance(choices, list) or not choices or not isinstance(choices[0], dict) or choices[0].get("finish_reason") != "stop":
        raise ValueError("Missing complete answer; inspect response.json and finish_reason")
    content = choices[0].get("message", {}).get("content")
    candidate = json.loads(content) if isinstance(content, str) else None
    if not isinstance(candidate, dict):
        raise ValueError("Model answer is not the requested JSON object")
    source = candidate.get("code")
    acceptance = code_acceptance(source)
    (out / "candidate.py").write_text(source, encoding="utf-8")
    record = {"status": "completed", "task": "coding-task", "requested_model": model, "returned_model": response.get("model"), "request_id": response.get("id"), "elapsed_ms": round((time.perf_counter() - started) * 1000), "usage": response.get("usage"), "billed_usd": None, "acceptance": acceptance}
    save_json(out / "record.json", redact(record, key))
    return record


def execute_edit(model, image_path, key, out, request=request_json, download=read_response):
    image = image_path.read_bytes()
    if len(image) > 10 * 1024 * 1024:
        raise ValueError("PNG input exceeds this example's 10 MiB limit")
    dimensions = png_dimensions(image)
    body, content_type = multipart(image, model, "quickrouter-" + uuid.uuid4().hex)
    prepare_out(out)
    (out / "input.png").write_bytes(image)
    save_json(out / "request.json", {"model": model, "prompt": EDIT_PROMPT, "n": 1, "size": "1024x1024", "quality": "low", "input_sha256": hashlib.sha256(image).hexdigest(), "input_dimensions": dimensions})
    started = time.perf_counter()
    response = request(API + "/images/edits", body, key, content_type)
    save_json(out / "response.json", redact(response, key))
    values = response.get("data")
    if not isinstance(values, list) or not values or not isinstance(values[0], dict):
        raise ValueError("Missing output image; inspect response.json")
    item = values[0]
    if isinstance(item.get("b64_json"), str):
        result = base64.b64decode(item["b64_json"], validate=True)
    elif isinstance(item.get("url"), str):
        result = download(Request(safe_url(item["url"]), headers={"Accept": "image/*"}), 32 * 1024 * 1024)
    else:
        raise ValueError("Expected data[0].b64_json or a HTTPS image URL")
    extension = image_extension(result)
    filename = "edited." + extension
    (out / filename).write_bytes(result)
    record = {"status": "completed", "task": "image-editing", "requested_model": model, "elapsed_ms": round((time.perf_counter() - started) * 1000), "usage": response.get("usage"), "billed_usd": None, "file": filename, "output_sha256": hashlib.sha256(result).hexdigest(), "manual_acceptance_required": ["Product color, shape, label, quantity and proportions", "Background and lighting", "Open the full image and inspect dimensions"]}
    save_json(out / "record.json", redact(record, key))
    return record


def download_video(result, out, download=read_response):
    if not isinstance(result, dict) or result.get("status") != "succeeded" or not isinstance(result.get("id"), str):
        raise ValueError("Expected a succeeded result with task ID from seedance-task.py")
    request = Request(safe_url(result.get("video_url")), headers={"Accept": "video/mp4"})
    out.mkdir(parents=True, exist_ok=False)
    data = download(request, 256 * 1024 * 1024)
    if len(data) < 16 or data[4:8] != b"ftyp":
        raise ValueError("Downloaded content is not an MP4 container; check the result URL")
    (out / "result.mp4").write_bytes(data)
    record = {"status": "downloaded", "task_id": result["id"], "bytes": len(data), "sha256": hashlib.sha256(data).hexdigest(), "usage": result.get("usage"), "billed_usd": None, "validation_scope": "MP4 container signature only; inspect playback, duration and visual content manually"}
    save_json(out / "record.json", record)
    return record


def main():
    parser = argparse.ArgumentParser(description="QuickRouter runnable examples; no request without --run. Coding/edit requests may be charged.")
    parser.add_argument("action", choices=["code", "edit", "download-video"])
    parser.add_argument("--model")
    parser.add_argument("--image", type=Path)
    parser.add_argument("--result", type=Path)
    parser.add_argument("--out", type=Path)
    parser.add_argument("--run", action="store_true")
    args = parser.parse_args()
    model = args.model or ("gpt-image-2" if args.action == "edit" else "glm-5.3")
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._/-]{0,99}", model):
        parser.error("Provide an exact valid model ID")
    if args.action == "edit" and not args.image:
        parser.error("edit requires --image input.png")
    if args.action == "download-video" and not args.result:
        parser.error("download-video requires --result video-result.json")
    if not args.run:
        preview = {"mode": "preview", "network_request_sent": False, "action": args.action}
        if args.action == "code":
            preview.update({"url": API + "/chat/completions", "body": code_payload(model), "fixture": FIXTURE, "expected_cases": CASES})
        elif args.action == "edit":
            preview.update({"url": API + "/images/edits", "model": model, "input": str(args.image), "prompt": EDIT_PROMPT, "size": "1024x1024", "quality": "low"})
        else:
            preview.update({"result": str(args.result), "authorization_header_sent": False})
        print(json.dumps(preview, ensure_ascii=False, indent=2))
        return 0
    if not args.out:
        parser.error("--run requires --out with a NEW output directory")
    if args.action == "download-video":
        result = download_video(json.loads(args.result.read_text(encoding="utf-8")), args.out)
    else:
        key = os.environ.get("QUICKROUTER_API_KEY", "")
        if not key.strip():
            parser.error("Set QUICKROUTER_API_KEY locally; do not put the Key into this file")
        result = execute_code(model, key, args.out) if args.action == "code" else execute_edit(model, args.image, key, args.out)
    print(json.dumps(redact(result, os.environ.get("QUICKROUTER_API_KEY", "")), ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (OSError, ValueError, SyntaxError, TypeError) as error:
        print(redact(str(error), os.environ.get("QUICKROUTER_API_KEY", "")), file=sys.stderr)
        sys.exit(1)