#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import urllib.error
import urllib.request
from typing import Any

DEFAULT_GATEWAY_URL = os.environ.get("SOC_MEMORY_GATEWAY_URL", "http://127.0.0.1:1934")
DEFAULT_GATEWAY_API_KEY = os.environ.get("SOC_MEMORY_GATEWAY_API_KEY", "")

URI_PREFIXES = {
    "case": "viking://resources/soc-memory-poc/case",
    "knowledge": "viking://resources/soc-memory-poc/knowledge",
    "all": "viking://resources/soc-memory-poc",
}


def post_json(url: str, payload: dict[str, Any], api_key: str = "") -> dict[str, Any]:
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, method="POST")
    req.add_header("Content-Type", "application/json")
    if api_key:
        req.add_header("X-API-Key", api_key)
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode("utf-8"))


def canonicalize_uri(uri: str) -> str:
    if ".json/" in uri:
        return uri.split(".json/", 1)[0] + ".json"
    return uri


def filter_results(results: list[dict[str, Any]], prefix: str) -> list[dict[str, Any]]:
    deduped: dict[str, dict[str, Any]] = {}
    for item in results:
        uri = item.get("uri") or ""
        canonical = canonicalize_uri(uri)
        if not canonical.startswith(prefix):
            continue
        score = item.get("score") or 0
        payload = dict(item)
        payload["uri"] = canonical
        if canonical not in deduped or score > (deduped[canonical].get("score") or 0):
            deduped[canonical] = payload
    return sorted(deduped.values(), key=lambda entry: entry.get("score") or 0, reverse=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Search SOC Memory Gateway for case / knowledge context.")
    parser.add_argument("--query", required=True, help="Search query")
    parser.add_argument("--kind", choices=["case", "knowledge", "all"], default="all", help="SOC resource scope")
    parser.add_argument("--limit", type=int, default=5, help="Max results")
    parser.add_argument("--gateway-url", default=DEFAULT_GATEWAY_URL, help="Memory Gateway base URL")
    parser.add_argument("--api-key", default=DEFAULT_GATEWAY_API_KEY, help="Gateway API key if required")
    args = parser.parse_args()

    prefix = URI_PREFIXES[args.kind]
    payload = {
        "query": args.query,
        "limit": max(args.limit * 5, 10),
        "uri": prefix,
    }
    try:
        result = post_json(args.gateway_url.rstrip("/") + "/api/search", payload, api_key=args.api_key)
    except urllib.error.URLError as exc:
        raise SystemExit(f"Gateway search failed: {exc}") from exc

    raw_results = result.get("results", [])
    filtered = filter_results(raw_results, prefix)
    output = {
        "query": args.query,
        "kind": args.kind,
        "uri_prefix": prefix,
        "results": filtered[: args.limit],
        "total": len(filtered),
    }
    print(json.dumps(output, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()