#!/usr/bin/env python3
"""Build source-specific canonical claim context from a Postgres dump without a database write."""

from __future__ import annotations

import argparse
import codecs
import hashlib
import importlib.util
import json
import os
import re
import shutil
import subprocess
import tempfile
from collections import Counter
from decimal import Decimal, InvalidOperation
from pathlib import Path
from types import ModuleType
from typing import Any

SCHEMA = "livingip.kbSourceDumpContextReceipt.v1"
STATUS_SCHEMA = "livingip.kbSourceDumpContextStatus.v1"
DEFAULT_DOCKER_IMAGE = "postgres:16-alpine"
DEFAULT_LIMIT = 12
DEFAULT_COPY_ENCODING = "utf-8"
MAX_PREVIEW_BYTES = 16_000
TABLES = ("claims", "claim_evidence", "claim_edges")


class ContextError(RuntimeError):
    """Raised when a dump cannot safely produce canonical retrieval context."""


def sha256_bytes(value: bytes) -> str:
    return hashlib.sha256(value).hexdigest()


def _read_file(path: Path, label: str) -> bytes:
    try:
        if not path.is_file():
            raise ContextError(f"{label} is not a regular file: {path}")
        return path.read_bytes()
    except OSError as exc:
        raise ContextError(f"could not read {label}: {exc}") from exc


def _load_kb_tool(repo_root: Path) -> ModuleType:
    path = repo_root / "hermes-agent" / "leoclean-bin" / "kb_tool.py"
    spec = importlib.util.spec_from_file_location("teleo_kb_tool_context", path)
    if spec is None or spec.loader is None:
        raise ContextError(f"could not load retrieval semantics from {path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _copy_encoding(value: str) -> str:
    try:
        return codecs.lookup(value).name
    except LookupError as exc:
        raise ContextError("COPY encoding is not supported") from exc


def _decode_copy_field(value: str, *, encoding: str = DEFAULT_COPY_ENCODING) -> str | None:
    if value == r"\N":
        return None
    encoding = _copy_encoding(encoding)
    result = bytearray()
    index = 0
    escapes = {"b": 8, "f": 12, "n": 10, "r": 13, "t": 9, "v": 11, "\\": 92}

    def append_text(text: str) -> None:
        try:
            result.extend(text.encode(encoding, errors="strict"))
        except UnicodeEncodeError as exc:
            raise ContextError(f"COPY field cannot be encoded as {encoding}") from exc

    while index < len(value):
        char = value[index]
        if char != "\\":
            append_text(char)
            index += 1
            continue
        if index + 1 >= len(value):
            raise ContextError("COPY field contains a dangling backslash escape")
        index += 1
        escaped = value[index]
        if escaped in escapes:
            result.append(escapes[escaped])
            index += 1
            continue
        if escaped == "x":
            match = re.match(r"[0-9A-Fa-f]{1,2}", value[index + 1 :])
            if match is not None:
                result.append(int(match.group(0), 16))
                index += 1 + len(match.group(0))
                continue
        if escaped in "01234567":
            match = re.match(r"[0-7]{1,3}", value[index:])
            if match is None:
                raise ContextError("COPY field contains a malformed octal byte escape")
            byte_value = int(match.group(0), 8)
            if byte_value > 0xFF:
                raise ContextError("COPY field contains an out-of-range octal byte escape")
            result.append(byte_value)
            index += len(match.group(0))
            continue
        append_text(escaped)
        index += 1
    if 0 in result:
        raise ContextError("COPY field contains a zero byte")
    try:
        return result.decode(encoding, errors="strict")
    except UnicodeDecodeError as exc:
        raise ContextError(f"COPY byte escapes do not form valid {encoding} text") from exc


def parse_copy_table(
    sql: str, table: str, *, encoding: str = DEFAULT_COPY_ENCODING
) -> list[dict[str, str | None]]:
    encoding = _copy_encoding(encoding)
    prefix = f"COPY public.{table} ("
    lines = iter(line[:-1] if line.endswith("\r") else line for line in sql.split("\n"))
    for line in lines:
        if line.startswith(prefix) and line.endswith(") FROM stdin;"):
            columns = [item.strip() for item in line[len(prefix) : -len(") FROM stdin;")].split(",")]
            break
    else:
        raise ContextError(f"pg_restore output has no COPY block for public.{table}")
    rows: list[dict[str, str | None]] = []
    for line in lines:
        if line == r"\.":
            return rows
        fields = line.split("\t")
        if len(fields) != len(columns):
            raise ContextError(f"public.{table} COPY row has {len(fields)} fields; expected {len(columns)}")
        rows.append(
            dict(zip(columns, (_decode_copy_field(field, encoding=encoding) for field in fields), strict=True))
        )
    raise ContextError(f"public.{table} COPY block is unterminated")


def restore_table(
    dump_path: Path, table: str, docker_image: str, *, encoding: str = DEFAULT_COPY_ENCODING
) -> str:
    if table not in TABLES:
        raise ContextError(f"unsupported canonical table: {table}")
    if shutil.which("docker") is None:
        raise ContextError("docker is not installed or not available on PATH")
    command = [
        "docker",
        "run",
        "--rm",
        "--network=none",
        "--mount",
        f"type=bind,src={dump_path.resolve()},dst=/canonical.dump,readonly",
        docker_image,
        "pg_restore",
        "--data-only",
        "-t",
        table,
        "-f",
        "-",
        "/canonical.dump",
    ]
    encoding = _copy_encoding(encoding)
    try:
        completed = subprocess.run(
            command,
            capture_output=True,
            text=True,
            encoding=encoding,
            errors="strict",
            timeout=180,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired, UnicodeError) as exc:
        raise ContextError(f"pg_restore failed ({type(exc).__name__}); diagnostic detail redacted") from exc
    if completed.returncode != 0:
        raise ContextError(f"pg_restore returned exit {completed.returncode}; diagnostic detail redacted")
    return completed.stdout


def _decimal(value: str | None) -> Decimal:
    if value is None:
        return Decimal(0)
    try:
        return Decimal(value)
    except InvalidOperation as exc:
        raise ContextError("canonical claim confidence is not numeric") from exc


def build_context(
    *,
    dump_path: Path,
    artifact_path: Path,
    title: str,
    limit: int,
    docker_image: str,
    repo_root: Path,
    copy_encoding: str = DEFAULT_COPY_ENCODING,
) -> tuple[dict[str, Any], dict[str, Any]]:
    if limit < 1 or limit > 20:
        raise ContextError("limit must be between 1 and 20")
    dump_bytes = _read_file(dump_path, "canonical dump")
    artifact_bytes = _read_file(artifact_path, "source artifact")
    try:
        preview = artifact_bytes[:MAX_PREVIEW_BYTES].decode("utf-8", errors="strict")
    except UnicodeDecodeError as exc:
        raise ContextError(f"source artifact preview must be strict UTF-8: {exc}") from exc
    kb_tool = _load_kb_tool(repo_root)
    query = " ".join(kb_tool.query_terms(f"{title} {preview}"))
    terms = kb_tool.query_terms(query)
    min_score = 2 if len(terms) >= 3 else 1

    copy_encoding = _copy_encoding(copy_encoding)
    tables = {
        table: parse_copy_table(
            restore_table(dump_path, table, docker_image, encoding=copy_encoding),
            table,
            encoding=copy_encoding,
        )
        for table in TABLES
    }
    evidence_counts = Counter(row["claim_id"] for row in tables["claim_evidence"] if row["claim_id"])
    edge_counts: Counter[str] = Counter()
    for row in tables["claim_edges"]:
        if row["from_claim"]:
            edge_counts[row["from_claim"]] += 1
        if row["to_claim"]:
            edge_counts[row["to_claim"]] += 1

    ranked: list[dict[str, Any]] = []
    for row in tables["claims"]:
        if row["status"] != "open" or not row["id"] or not row["text"]:
            continue
        text = row["text"]
        searchable = f"{text}\n{row['tags'] or ''}".lower()
        score = sum(term in searchable for term in terms)
        if score < min_score:
            continue
        ranked.append(
            {
                "id": row["id"],
                "text": text,
                "score": score,
                "evidence_count": evidence_counts[row["id"]],
                "edge_count": edge_counts[row["id"]],
                "_confidence": _decimal(row["confidence"]),
            }
        )
    ranked.sort(
        key=lambda row: (
            -row["score"],
            -row["evidence_count"],
            -row["edge_count"],
            -row["_confidence"],
            len(row["text"]),
            row["text"],
            row["id"],
        )
    )
    candidates = [{key: value for key, value in row.items() if key != "_confidence"} for row in ranked[:limit]]
    context = {"database_search_query": query, "candidate_claims": candidates}
    receipt = {
        "schema": SCHEMA,
        "status": "prepared",
        "database_write_performed": False,
        "model_call_performed": False,
        "canonical_dump": {
            "path": str(dump_path.resolve()),
            "sha256": sha256_bytes(dump_bytes),
        },
        "source": {
            "path": str(artifact_path.resolve()),
            "sha256": sha256_bytes(artifact_bytes),
        },
        "retrieval": {
            "query_sha256": sha256_bytes(query.encode("utf-8")),
            "term_count": len(terms),
            "minimum_score": min_score,
            "canonical_open_claim_count": sum(row["status"] == "open" for row in tables["claims"]),
            "candidate_count": len(candidates),
            "candidate_ids": [row["id"] for row in candidates],
        },
        "runtime": {
            "docker_image": docker_image,
            "network_disabled": True,
            "pg_restore_output_retained": False,
            "copy_encoding": copy_encoding,
            "copy_decode_errors": "strict",
        },
    }
    return context, receipt


def _stage_private_json(path: Path, value: Any) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
    os.chmod(path.parent, 0o700)
    fd, raw_temp = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent)
    temp = Path(raw_temp)
    try:
        os.fchmod(fd, 0o600)
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            json.dump(value, handle, indent=2, sort_keys=True, ensure_ascii=True)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        return temp
    except BaseException:
        if temp.exists():
            temp.unlink()
        raise


def _private_json(path: Path, value: Any) -> None:
    if path.exists():
        raise ContextError(f"refusing to overwrite existing output: {path}")
    temp = _stage_private_json(path, value)
    try:
        os.replace(temp, path)
        os.chmod(path, 0o600)
    finally:
        if temp.exists():
            temp.unlink()


def _publish_private_json_outputs(outputs: tuple[tuple[Path, Any], ...]) -> None:
    for path, _value in outputs:
        if path.exists():
            raise ContextError(f"refusing to overwrite existing output: {path}")

    staged: list[tuple[Path, Path]] = []
    try:
        for path, value in outputs:
            staged.append((path, _stage_private_json(path, value)))
        for path, temp in staged:
            os.replace(temp, path)
            os.chmod(path, 0o600)
    finally:
        for _path, temp in staged:
            if temp.exists():
                temp.unlink()


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dump", required=True, type=Path)
    parser.add_argument("--artifact", required=True, type=Path)
    parser.add_argument("--title", required=True)
    parser.add_argument("--context-output", required=True, type=Path)
    parser.add_argument("--receipt-output", required=True, type=Path)
    parser.add_argument("--limit", type=int, default=DEFAULT_LIMIT)
    parser.add_argument("--docker-image", default=DEFAULT_DOCKER_IMAGE)
    parser.add_argument("--copy-encoding", default=DEFAULT_COPY_ENCODING)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    repo_root = Path(__file__).resolve().parents[1]
    context_existed = args.context_output.is_file()
    receipt_existed = args.receipt_output.is_file()
    try:
        if args.context_output.resolve() == args.receipt_output.resolve():
            raise ContextError("context and receipt outputs must be different files")
        context, receipt = build_context(
            dump_path=args.dump,
            artifact_path=args.artifact,
            title=args.title,
            limit=args.limit,
            docker_image=args.docker_image,
            repo_root=repo_root,
            copy_encoding=args.copy_encoding,
        )
        receipt["outputs"] = {
            "context": str(args.context_output.resolve()),
            "receipt": str(args.receipt_output.resolve()),
        }
        _publish_private_json_outputs(
            ((args.context_output, context), (args.receipt_output, receipt))
        )
    except (ContextError, OSError, ValueError):
        context_written = not context_existed and args.context_output.is_file()
        receipt_written = not receipt_existed and args.receipt_output.is_file()
        print(
            json.dumps(
                {
                    "schema": STATUS_SCHEMA,
                    "status": "rejected",
                    "database_write_performed": False,
                    "model_call_performed": False,
                    "private_context_written": context_written,
                    "private_receipt_written": receipt_written,
                    "reason": "context_preparation_failed",
                },
                sort_keys=True,
            )
        )
        return 2
    print(
        json.dumps(
            {
                "schema": STATUS_SCHEMA,
                "status": "prepared",
                "database_write_performed": False,
                "model_call_performed": False,
                "private_context_written": True,
                "private_receipt_written": True,
            },
            indent=2,
            sort_keys=True,
        )
    )
    return 0


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