"""Small, dependency-free client shared by Claude skills and the MCP bridge."""

from __future__ import annotations

import argparse
from dataclasses import dataclass
import ipaddress
import json
import os
import re
from pathlib import Path
import stat
import urllib.error
import urllib.parse
import urllib.request

MAX_RESPONSE_BYTES = 2 * 1024 * 1024
ENGINES = ("rag", "graphrag")
STATUSES = {"ok", "no_answer", "ambiguous", "error"}


class ClientError(Exception):
    """The message is a public reason code, never an exception from the transport."""


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


@dataclass(frozen=True)
class Config:
    url: str
    key: str
    timeout: float = 300
    allow_http: bool = False

    @classmethod
    def load(cls):
        values = {}
        filename = os.getenv("ZP_AGENT_CONFIG", "").strip()
        if filename:
            try:
                path = Path(filename).expanduser()
                if os.name == "posix" and stat.S_IMODE(path.stat().st_mode) & 0o077:
                    raise ClientError("config_permissions_require_0600")
                values = json.loads(path.read_text())
                if not isinstance(values, dict):
                    raise ValueError()
            except ClientError:
                raise
            except (OSError, ValueError):
                raise ClientError("invalid_configuration") from None
        try:
            config = cls(
                os.getenv("ZP_AGENT_URL", values.get("backend_url", "")).strip().rstrip("/"),
                os.getenv("ZP_AGENT_API_KEY", values.get("api_key", "")).strip(),
                float(os.getenv("ZP_AGENT_TIMEOUT", values.get("timeout_seconds", 300))),
                os.getenv("ZP_AGENT_ALLOW_HTTP", "").lower() in {"1", "true"},
            )
            config.validate()
            return config
        except (ValueError, TypeError, AttributeError):
            raise ClientError("invalid_configuration") from None

    def validate(self):
        try:
            url = urllib.parse.urlsplit(self.url)
            if (
                url.scheme not in {"http", "https"}
                or not url.hostname
                or url.username is not None
                or url.password is not None
                or url.query
                or url.fragment
                or any(c.isspace() for c in self.url)
            ):
                raise ValueError()
            url.port
            if url.scheme == "http":
                local = url.hostname == "localhost"
                try:
                    local = local or ipaddress.ip_address(url.hostname).is_loopback
                except ValueError:
                    pass
                if not local and not self.allow_http:
                    raise ValueError()
            if not self.key or "\r" in self.key or "\n" in self.key or not 1 <= self.timeout <= 600:
                raise ValueError()
        except (ValueError, TypeError):
            raise ClientError("invalid_configuration") from None


def error_result(reason):
    messages = {
        "reference_unavailable": "Zdrojový repozitár sa nepodarilo verejne overiť. Odpoveď založenú iba na tomto odkaze neuvádzam.",
        "unauthorized": "Prístup nebol autorizovaný. Skontroluj prihlásenie konektora a jeho konfiguráciu.",
        "forbidden": "Pre tento nástroj nemáš potrebné oprávnenie.",
        "invalid_query": "Otázku sa nepodarilo spracovať. Zadaj konkrétnu otázku a primeraný počet výsledkov.",
        "question_mismatch": "Odpoveď patrí k inej otázke a nezobrazí sa. Zopakuj aktuálnu otázku.",
        "engine_mismatch": "Príkaz a zvolený nástroj sa nezhodujú. Použi nástroj podľa /rag alebo /graphrag.",
        "model_unavailable": "Jazykový model momentálne nie je dostupný. Skús otázku neskôr.",
        "answer_rejected": "Odpoveď sa nepodarilo spoľahlivo overiť podľa podkladov. "
        "Neoverené údaje sa nezobrazia.",
        "answer_unavailable": "Overenú odpoveď sa momentálne nepodarilo pripraviť. Skús to neskôr.",
        "answer_capacity_exceeded": "Agent práve spracúva iné otázky. Skús to o chvíľu.",
        "backend_busy": "Agent práve spracúva iné otázky. Skús to o chvíľu.",
        "backend_timeout": "Vyhľadávanie neodpovedalo včas. Skús otázku zopakovať neskôr.",
        "backend_unavailable": "Spojenie s backendom momentálne nie je dostupné. Skontroluj, či bežia služby a tunel.",
        "retrieval_unavailable": "Vyhľadávací index momentálne nie je dostupný. Skús otázku neskôr.",
        "content_index_unavailable": "Obsahové vyhľadávanie momentálne nie je pripravené alebo sa jeho výstup "
        "nepodarilo overiť. Treba overiť extrakciu pre aktuálnu verziu dát; "
        "nejde o potvrdenie, že hľadané údaje neexistujú.",
    }
    return {
        "status": "error",
        "reason": reason,
        "accepted_for_delivery": False,
        "answer": messages.get(reason, "Overenú odpoveď sa momentálne nepodarilo získať."),
        "sources": [],
        "index_version": None,
        "complete": False,
        "truncated": False,
        "truncation_reasons": [],
    }


def request_json(config: Config, endpoint: str, body=None):
    config.validate()
    request = urllib.request.Request(
        config.url + endpoint,
        data=json.dumps(body, ensure_ascii=False).encode() if body is not None else None,
        headers={
            "Authorization": "Bearer " + config.key,
            "Accept": "application/json",
            **({"Content-Type": "application/json"} if body is not None else {}),
        },
    )
    try:
        with urllib.request.build_opener(NoRedirect()).open(
            request, timeout=config.timeout
        ) as response:
            raw = response.read(MAX_RESPONSE_BYTES + 1)
    except urllib.error.HTTPError as exc:
        code = exc.code
        exc.close()
        reason = {
            401: "unauthorized",
            403: "forbidden",
            404: "backend_endpoint_not_found",
            400: "invalid_query",
            422: "invalid_request",
            429: "backend_busy",
            503: "backend_unavailable",
        }.get(code, "backend_http_error")
        raise ClientError("redirect_not_allowed" if 300 <= code < 400 else reason) from None
    except (urllib.error.URLError, TimeoutError, OSError) as exc:
        timeout = isinstance(exc, TimeoutError) or isinstance(
            getattr(exc, "reason", None), TimeoutError
        )
        raise ClientError("backend_timeout" if timeout else "backend_unavailable") from None
    if len(raw) > MAX_RESPONSE_BYTES:
        raise ClientError("response_too_large")
    try:
        value = json.loads(raw)
        if not isinstance(value, dict):
            raise ValueError()
        return value
    except (ValueError, UnicodeError):
        raise ClientError("invalid_backend_response") from None


def validate_metadata(value):
    if (
        not {"complete", "truncated", "truncation_reasons", "index_version"} <= value.keys()
        or value.get("complete") is not None
        and type(value["complete"]) is not bool
        or type(value.get("truncated")) is not bool
        or value["truncated"]
        and value.get("complete") is not False
        or not isinstance(value.get("truncation_reasons"), list)
        or not all(isinstance(x, str) for x in value["truncation_reasons"])
        or value["truncation_reasons"]
        and not value["truncated"]
        or value.get("index_version") is not None
        and not isinstance(value["index_version"], dict)
    ):
        raise ClientError("invalid_backend_response")


def query_backend(query: str, engine: str, *, mode="answer", limit=5, config=None):
    """Return only approved answer fields, or explicitly marked retrieval context."""
    try:
        if (
            engine not in ENGINES
            or mode not in {"answer", "retrieve"}
            or type(limit) is not int
            or not 1 <= limit <= 20
            or not isinstance(query, str)
        ):
            raise ClientError("invalid_query")
        query = query.strip()
        command = re.match(r"^/(rag|graphrag)(?=\s|$)", query, re.I)
        if command:
            if command[1].lower() != engine:
                raise ClientError("engine_mismatch")
            query = query[command.end() :].strip()
        if not 1 <= len(query) <= 500:
            raise ClientError("invalid_query")
        config = config or Config.load()
        endpoint = "/answer" if mode == "answer" else "/" + engine
        body = {"query": query.strip(), "limit": limit}
        if mode == "answer" or engine == "rag":
            body["include_diagnostics"] = False
        if mode == "answer":
            body["engine"] = engine
        else:
            body["published_only"] = True
        value = request_json(config, endpoint, body)
        if not isinstance(value.get("status"), str) or value["status"] not in STATUSES:
            raise ClientError("invalid_backend_response")
        if value["status"] == "error":
            # Raw model output, diagnostics and rejected text must not escape.
            reason = value.get("reason")
            safe_reasons = {
                "model_unavailable",
                "answer_capacity_exceeded",
                "answer_rejected",
                "answer_unavailable",
                "retrieval_unavailable",
                "content_index_unavailable",
                "reference_unavailable",
            }
            raise ClientError(
                reason
                if isinstance(reason, str) and reason in safe_reasons
                else "backend_answer_error" if mode == "answer" else "retrieval_unavailable"
            )
        validate_metadata(value)
        fields = {
            "status",
            "reason",
            "index_version",
            "complete",
            "truncated",
            "truncation_reasons",
        }
        if mode == "answer":
            if not isinstance(value.get("query"), str):
                raise ClientError("invalid_backend_response")
            if value["query"] != body["query"] or value.get("engine") != engine:
                raise ClientError("question_mismatch")
            if value.get("accepted_for_delivery") is not True:
                raise ClientError("answer_rejected")
            if (
                not isinstance(value.get("answer"), str)
                or not value["answer"].strip()
                or not isinstance(value.get("source_urls"), list)
                or not all(isinstance(x, str) for x in value["source_urls"])
                or not isinstance(value.get("claims"), list)
            ):
                raise ClientError("invalid_backend_response")
            if "answer_coverage" in value:
                coverage = value["answer_coverage"]
                if (
                    not isinstance(coverage, dict)
                    or coverage.get("scope")
                    not in ("fact", "overview", "assignment", "list", "count")
                    or coverage.get("status")
                    not in ("supported", "partial", "unknown", "not_found", "unavailable")
                    or any(
                        not isinstance(coverage.get(key), list)
                        or any(not isinstance(item, str) for item in coverage[key])
                        for key in ["warnings", "limitations"]
                    )
                ):
                    raise ClientError("invalid_backend_response")
            if "knowledge_version" in value and not isinstance(value["knowledge_version"], str):
                raise ClientError("invalid_backend_response")
            fields |= {
                "query",
                "answer",
                "source_urls",
                "claims",
                "count",
                "accepted_for_delivery",
                "answer_coverage",
            }
            fields.add("knowledge_version")
        else:
            if not isinstance(value.get("sources"), list):
                raise ClientError("invalid_backend_response")
            if engine == "rag" and not isinstance(value.get("context"), str):
                raise ClientError("invalid_backend_response")
            if engine == "graphrag" and (
                not isinstance(value.get("facts"), list) or not isinstance(value.get("paths"), list)
            ):
                raise ClientError("invalid_backend_response")
            fields |= {
                "sources",
                "context",
                "facts",
                "supporting_facts",
                "paths",
                "aggregation",
                "constraints",
                "answer_candidates",
                "instructions",
            }
        result = {key: item for key, item in value.items() if key in fields}
        return {
            **result,
            "engine": engine,
            "output_kind": "verified_answer" if mode == "answer" else "retrieval_context",
            **({"requires_final_answer_verification": True} if mode == "retrieve" else {}),
        }
    except ClientError as exc:
        return error_result(str(exc))


def backend_health(config=None):
    try:
        value = request_json(config or Config.load(), "/health")
        tools = value.get("tools")
        if (
            type(value.get("ready")) is not bool
            or not isinstance(tools, dict)
            or any(
                not isinstance(tools.get(engine), dict)
                or type(tools[engine].get("ready")) is not bool
                for engine in ENGINES
            )
        ):
            raise ClientError("invalid_backend_response")
        return {
            "status": "ok" if value["ready"] else "error",
            "ready": value["ready"],
            "index_version": value.get("index_version"),
            "tools": {
                engine: {
                    k: v
                    for k, v in tools[engine].items()
                    if k in {"ready", "reason", "index_version"}
                }
                for engine in ENGINES
            },
        }
    except ClientError as exc:
        return {**error_result(str(exc)), "ready": False}


def answer_for_presentation(result: dict) -> dict:
    """Keep approved prose/caveats intact; retrieval telemetry is not answer prose.

    Called after query_backend validates the complete backend contract. Full
    metadata remains available there, in /answer and in diagnostic retrieval.
    Never change an incomplete result into a claim of completeness.
    """
    fields = {
        "status",
        "reason",
        "query",
        "engine",
        "output_kind",
        "accepted_for_delivery",
        "answer",
        "source_urls",
        "index_version",
        "knowledge_version",
    }
    return {key: value for key, value in result.items() if key in fields}


def main(engine=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("query", nargs="?")
    if engine is None:
        parser.add_argument("--engine", choices=ENGINES, default="rag")
    parser.add_argument("--mode", choices=["answer", "retrieve"], default="answer")
    parser.add_argument("--limit", type=int, default=5)
    parser.add_argument("--health", action="store_true")
    args = parser.parse_args()
    result = (
        backend_health()
        if args.health
        else query_backend(args.query, engine or args.engine, mode=args.mode, limit=args.limit)
    )
    if not args.health and args.mode == "answer":
        result = answer_for_presentation(result)
    print(json.dumps(result, ensure_ascii=False, separators=(",", ":")))
    return 1 if result.get("status") == "error" else 0


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