"""Ask the local LightRAG server a question that exists in the current graph.

Usage:
    python ask_aclasta.py
"""

from __future__ import annotations

import argparse
import json
import sys
import urllib.error
import urllib.request
from pathlib import Path

from pharmacy_agent import QUERY_SETTINGS, prepare_query


DEFAULT_URL = "http://localhost:9621"
DEFAULT_MODE = "mix"
DEFAULT_QUESTION = "Ake su kontraindikacie lieku ACARIZAX?"


if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8")


def post_query(base_url: str, question: str, mode: str, timeout: int) -> dict:
    index = Path(__file__).resolve().parent / "outputs/demo/product_index.json"
    products = json.loads(index.read_text(encoding="utf-8")) if index.exists() else []
    prepared = prepare_query(question, products)
    payload = json.dumps(
        {**QUERY_SETTINGS, "query": prepared["query"], "user_prompt": prepared["user_prompt"],
         "mode": mode, "include_references": True, "include_chunk_content": True},
        ensure_ascii=False,
    ).encode("utf-8")
    request = urllib.request.Request(
        f"{base_url.rstrip('/')}/query",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def main() -> None:
    parser = argparse.ArgumentParser(description="Query LightRAG about ACARIZAX contraindications.")
    parser.add_argument("--url", default=DEFAULT_URL, help="LightRAG server URL.")
    parser.add_argument("--mode", default=DEFAULT_MODE, choices=("mix", "hybrid", "local", "global", "naive"))
    parser.add_argument("--timeout", type=int, default=120)
    parser.add_argument("--question", default=DEFAULT_QUESTION)
    args = parser.parse_args()

    try:
        result = post_query(args.url, args.question, args.mode, args.timeout)
    except urllib.error.URLError as exc:
        raise SystemExit(
            "LightRAG server is not reachable. Start it first with:\n"
            "  python start_servers.py\n\n"
            f"Details: {exc}"
        ) from exc
    except json.JSONDecodeError as exc:
        raise SystemExit(f"LightRAG returned invalid JSON: {exc}") from exc

    response = result.get("response")
    if response is None:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return

    print(response)
    references = result.get("references") or []
    if references:
        print("\nZdroje:")
        for ref in references:
            print(f"[{ref.get('reference_id', '?')}] {ref.get('file_path', '')}")


if __name__ == "__main__":
    main()