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


def main():
    parser = argparse.ArgumentParser(description="Chat with GLM over an OpenAI-compatible API")
    parser.add_argument("--prompt", required=True, help="User prompt")
    parser.add_argument("--system", default="You are a helpful assistant.", help="Optional system prompt")
    parser.add_argument("--model", default=os.getenv("GLM_PROXY_MODEL", "glm5-bd"), help="Model name")
    parser.add_argument("--json", action="store_true", help="Print raw JSON response")
    args = parser.parse_args()

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

    base_url = os.getenv("GLM_PROXY_BASE_URL", "https://api-ai.7e.ink/v1").rstrip("/")
    url = f"{base_url}/chat/completions"

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

    req = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        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:
        print(parsed["choices"][0]["message"]["content"])
    except Exception:
        print(json.dumps(parsed, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()