#!/usr/bin/env python3
"""Native web grounding wrapper via ZenMux chat completions.

Use this when you need reproducible, model-native web search (grounding) and
machine-readable citations.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from scripts.lib.zenmux_client import ZenMuxClient, load_secrets


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Grounded web query via ZenMux")
    parser.add_argument("query", help="Question or search prompt")
    parser.add_argument("--model", default="google/gemini-3.1-flash-lite-preview")
    parser.add_argument("--max-tokens", type=int, default=2400)
    parser.add_argument("--temperature", type=float, default=0.2)
    parser.add_argument("--json", action="store_true", help="Emit JSON envelope")
    parser.add_argument("--log-file", help="Optional JSONL call log path")
    parser.add_argument("--system", default=(
        "You are a research assistant. Use web grounding when helpful. "
        "Return concise facts with explicit source-backed statements."
    ))
    return parser


def main() -> int:
    args = build_parser().parse_args()
    load_secrets()
    log_file = Path(args.log_file) if args.log_file else None

    with ZenMuxClient(log_file=log_file) as client:
        result = client.chat_complete_with_meta(
            model=args.model,
            system=args.system,
            user=args.query,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            web_search=True,
            web_search_options={},
            tag="ground",
        )

    if args.json:
        payload = {
            "query": args.query,
            "model": args.model,
            "content": result["content"],
            "citations": result["citations"],
            "usage": result["usage"],
        }
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    else:
        print(result["content"])
        if result["citations"]:
            print("\nCitations:")
            for idx, url in enumerate(result["citations"], start=1):
                print(f"{idx}. {url}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())