#!/usr/bin/env python3
import argparse
import json
import os
import sys
import time
from typing import Any, Dict, Optional

import requests


DEFAULT_BASE_URL = "https://api.coze.cn"
DEFAULT_RUN_PATH = "/v1/workflow/run"
DEFAULT_HISTORY_PATH = "/v1/workflows/{workflow_id}/run_histories/{execute_id}"
DEFAULT_SAMPLE_PROMPT = (
    "请帮我排查以下亚马逊商品在英国及欧洲地区是否存在知识产权侵权风险。"
    "请直接给出明确结论（有侵权风险 或 未发现明显侵权风险）。"
)


def build_headers(token: str) -> Dict[str, str]:
    normalized = (token or "").strip()
    if normalized.lower().startswith("bearer "):
        normalized = normalized[7:].strip()
    return {
        "Authorization": f"Bearer {normalized}",
        "Content-Type": "application/json; charset=utf-8",
        "Accept-Charset": "utf-8",
    }


def join_url(base_url: str, path: str) -> str:
    return base_url.rstrip("/") + "/" + path.lstrip("/")


def print_block(title: str, payload: Any) -> None:
    print(f"\n=== {title} ===")
    if isinstance(payload, (dict, list)):
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    else:
        print(payload)


def request_json(method: str, url: str, headers: Dict[str, str], **kwargs: Any) -> Dict[str, Any]:
    started = time.time()
    response = requests.request(method, url, headers=headers, timeout=kwargs.pop("timeout", 60), **kwargs)
    elapsed_ms = int((time.time() - started) * 1000)
    body_text = response.text
    print_block(
        f"{method.upper()} {url} -> HTTP {response.status_code} ({elapsed_ms} ms)",
        body_text if len(body_text) < 5000 else body_text[:5000] + "\n...<truncated>...",
    )
    response.raise_for_status()
    try:
        return response.json()
    except Exception as exc:
        raise RuntimeError(f"response is not valid JSON: {exc}") from exc


def extract_execute_id(payload: Dict[str, Any]) -> str:
    for key in ("execute_id", "executeId"):
        value = payload.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    raise RuntimeError("execute_id not found in Coze response")


def payload_data_text(payload: Dict[str, Any]) -> str:
    data = payload.get("data")
    if isinstance(data, str):
        return data
    if isinstance(data, dict):
        inner = data.get("data")
        if isinstance(inner, str):
            return inner
    return ""


def workflow_status(payload: Dict[str, Any]) -> str:
    for key in ("status", "run_status", "workflow_status"):
        value = payload.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip().upper()
    data = payload.get("data")
    if isinstance(data, dict):
        for key in ("status", "run_status", "workflow_status"):
            value = data.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip().upper()
    return ""


def looks_finished(payload: Dict[str, Any]) -> bool:
    status = workflow_status(payload)
    return status in {"SUCCESS", "SUCCEEDED", "DONE", "FAILED", "ERROR", "CANCELLED"}


def build_sample_run_body(workflow_id: str, prompt: str) -> Dict[str, Any]:
    return {
        "workflow_id": workflow_id,
        "parameters": {
            "title_list": ["WIRESTER Green Aluminum Round Bike Bell"],
            "url_list": ["https://m.media-amazon.com/images/I/51dETWpSQXL._AC_SY450_.jpg"],
            "items": [
                {
                    "group_key": "diag::1",
                    "row_id": "1",
                    "asin": "B0D14L34S7",
                    "country": "英国",
                    "title": "WIRESTER Green Aluminum Round Bike Bell",
                    "url": "https://m.media-amazon.com/images/I/51dETWpSQXL._AC_SY450_.jpg",
                }
            ],
            "prompt": prompt,
        },
        "is_async": True,
    }


def poll_history(
    base_url: str,
    workflow_id: str,
    execute_id: str,
    token: str,
    history_path: str,
    poll_interval: float,
    max_polls: int,
) -> int:
    headers = build_headers(token)
    history_url = join_url(
        base_url,
        history_path.replace("{workflow_id}", workflow_id).replace("{execute_id}", execute_id),
    )
    for idx in range(1, max_polls + 1):
        print(f"\n--- poll #{idx} execute_id={execute_id} ---")
        payload = request_json("GET", history_url, headers)
        code = payload.get("code")
        status = workflow_status(payload)
        data_text = payload_data_text(payload)
        print(f"code={code!r} status={status or '-'} has_data={bool(data_text)}")
        if data_text:
            print_block("Resolved Data", data_text)
            return 0
        if looks_finished(payload):
            print("workflow reached terminal status but no data was returned")
            return 2
        time.sleep(max(0.2, poll_interval))
    print("history polling reached max attempts without terminal result")
    return 3


def main() -> int:
    parser = argparse.ArgumentParser(description="Diagnose Coze run/history connectivity from server.")
    parser.add_argument("--base-url", default=os.getenv("AIIMAGE_APPEARANCE_PATENT_COZE_BASE_URL", DEFAULT_BASE_URL))
    parser.add_argument("--workflow-id", default=os.getenv("AIIMAGE_APPEARANCE_PATENT_COZE_WORKFLOW_ID") or os.getenv("workflow_id"))
    parser.add_argument("--token", default=os.getenv("AIIMAGE_APPEARANCE_PATENT_COZE_TOKEN") or os.getenv("coze_token"))
    parser.add_argument("--run-path", default=os.getenv("AIIMAGE_APPEARANCE_PATENT_COZE_WORKFLOW_PATH", DEFAULT_RUN_PATH))
    parser.add_argument(
        "--history-path",
        default=os.getenv("AIIMAGE_APPEARANCE_PATENT_COZE_WORKFLOW_HISTORY_PATH", DEFAULT_HISTORY_PATH),
    )
    parser.add_argument("--execute-id", help="Only test history/poll for an existing execute_id.")
    parser.add_argument("--poll-interval", type=float, default=2.0)
    parser.add_argument("--max-polls", type=int, default=20)
    parser.add_argument("--prompt", default=DEFAULT_SAMPLE_PROMPT)
    parser.add_argument("--skip-run", action="store_true", help="Do not create a new run; only poll execute_id.")
    args = parser.parse_args()

    if not args.workflow_id:
        print("missing workflow id: pass --workflow-id or set AIIMAGE_APPEARANCE_PATENT_COZE_WORKFLOW_ID/workflow_id", file=sys.stderr)
        return 4
    if not args.token:
        print("missing token: pass --token or set AIIMAGE_APPEARANCE_PATENT_COZE_TOKEN/coze_token", file=sys.stderr)
        return 4

    print_block(
        "Config",
        {
            "base_url": args.base_url,
            "workflow_id": args.workflow_id,
            "run_path": args.run_path,
            "history_path": args.history_path,
            "has_token": bool(args.token),
            "execute_id": args.execute_id,
            "skip_run": args.skip_run,
        },
    )

    if args.execute_id:
        return poll_history(
            args.base_url,
            args.workflow_id,
            args.execute_id,
            args.token,
            args.history_path,
            args.poll_interval,
            args.max_polls,
        )

    if args.skip_run:
        print("--skip-run requires --execute-id", file=sys.stderr)
        return 4

    headers = build_headers(args.token)
    run_url = join_url(args.base_url, args.run_path)
    body = build_sample_run_body(args.workflow_id, args.prompt)
    print_block("Run Request Body", body)
    payload = request_json("POST", run_url, headers, json=body)
    execute_id = extract_execute_id(payload)
    print(f"\nexecute_id={execute_id}")
    return poll_history(
        args.base_url,
        args.workflow_id,
        execute_id,
        args.token,
        args.history_path,
        args.poll_interval,
        args.max_polls,
    )


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except requests.HTTPError as exc:
        response = exc.response
        if response is not None:
            print(f"HTTP error: {response.status_code} {response.text}", file=sys.stderr)
        else:
            print(f"HTTP error: {exc}", file=sys.stderr)
        raise
    except Exception as exc:
        print(f"fatal: {exc}", file=sys.stderr)
        raise