#!/usr/bin/env python3
import argparse
import json
import os
import sys
import urllib.request
import urllib.error


def main():
    parser = argparse.ArgumentParser(description="Search via Tongyi/Qwen with enable_search=true")
    parser.add_argument("--query", required=True, help="Search query")
    parser.add_argument("--model", default=os.getenv("TONGYI_SEARCH_MODEL", "qwen-plus"), help="Model name")
    parser.add_argument("--system", default="You are a helpful research assistant. Use web search results when relevant and answer clearly.", help="System prompt")
    parser.add_argument("--json", action="store_true", help="Print raw JSON response")
    args = parser.parse_args()

    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        print("Error: DASHSCOPE_API_KEY environment variable is not set.", file=sys.stderr)
        sys.exit(1)

    base_url = os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1").rstrip("/")
    url = f"{base_url}/chat/completions"

    payload = {
        "model": args.model,
        "messages": [
            {"role": "system", "content": args.system},
            {"role": "user", "content": args.query},
        ],
        "enable_search": True,
    }

    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=data,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            body = resp.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        err = e.read().decode("utf-8", errors="replace") if e.fp else ""
        print(f"HTTP {e.code}: {err}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Request failed: {e}", file=sys.stderr)
        sys.exit(1)

    parsed = json.loads(body)
    if args.json:
        print(json.dumps(parsed, ensure_ascii=False, indent=2))
        return

    try:
        text = parsed["choices"][0]["message"]["content"]
    except Exception:
        print(json.dumps(parsed, ensure_ascii=False, indent=2))
        return

    if isinstance(text, list):
        parts = []
        for item in text:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
            else:
                parts.append(str(item))
        print("\n".join(p for p in parts if p))
    else:
        print(text)


if __name__ == "__main__":
    main()