#!/usr/bin/env python3
"""Safe helpers for the isolated Hindsight vs Supermemory POC."""

from __future__ import annotations

import argparse
import hashlib
import json
import mimetypes
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


def request_json(
    method: str,
    url: str,
    *,
    token: str = "",
    payload: dict[str, Any] | None = None,
    timeout: float = 30.0,
) -> Any:
    headers = {"Accept": "application/json"}
    data = None
    if token:
        headers["Authorization"] = f"Bearer {token}"
    if payload is not None:
        headers["Content-Type"] = "application/json"
        data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            body = response.read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:300]
        raise RuntimeError(f"{method} {url}: HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"{method} {url}: {exc.reason}") from exc
    return json.loads(body) if body else {}


def require_token(env_name: str) -> str:
    token = os.environ.get(env_name, "").strip()
    if not token:
        raise RuntimeError(f"{env_name} is not set")
    return token


def request_multipart_json(
    url: str,
    *,
    token: str,
    fields: dict[str, str],
    file_path: Path,
    timeout: float,
) -> Any:
    boundary = f"supermemory-poc-{uuid.uuid4().hex}"
    chunks: list[bytes] = []
    for name, value in fields.items():
        chunks.extend(
            [
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
                value.encode("utf-8"),
                b"\r\n",
            ]
        )
    filename = file_path.name.replace('"', "_")
    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
    chunks.extend(
        [
            f"--{boundary}\r\n".encode(),
            f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'.encode(),
            f"Content-Type: {mime_type}\r\n\r\n".encode(),
            file_path.read_bytes(),
            b"\r\n",
            f"--{boundary}--\r\n".encode(),
        ]
    )
    req = urllib.request.Request(
        url,
        data=b"".join(chunks),
        headers={
            "Accept": "application/json",
            "Authorization": f"Bearer {token}",
            "Content-Type": f"multipart/form-data; boundary={boundary}",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            body = response.read()
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")[:300]
        raise RuntimeError(f"POST {url}: HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"POST {url}: {exc.reason}") from exc
    return json.loads(body) if body else {}


def sample_hindsight(args: argparse.Namespace) -> None:
    base = args.hindsight_url.rstrip("/")
    bank = urllib.parse.quote(args.bank_id, safe="")
    list_url = f"{base}/v1/default/banks/{bank}/memories/list"
    first_query = urllib.parse.urlencode(
        {"type": "observation", "state": "valid", "limit": 1, "offset": 0}
    )
    first = request_json("GET", f"{list_url}?{first_query}", timeout=args.timeout)
    total = int(first.get("total") or 0)
    if total < 1:
        raise RuntimeError("Hindsight returned no active observations")

    wanted = min(args.limit, total)
    page_size = min(args.page_size, wanted)
    page_count = max(1, (wanted + page_size - 1) // page_size)
    max_offset = max(0, total - page_size)
    offsets = (
        [0]
        if page_count == 1
        else [round(index * max_offset / (page_count - 1)) for index in range(page_count)]
    )

    rows: list[dict[str, Any]] = []
    seen: set[str] = set()
    imported_at = datetime.now(timezone.utc).isoformat()
    for offset in offsets:
        query = urllib.parse.urlencode(
            {
                "type": "observation",
                "state": "valid",
                "limit": page_size,
                "offset": offset,
            }
        )
        response = request_json("GET", f"{list_url}?{query}", timeout=args.timeout)
        for item in response.get("items") or []:
            item_id = str(item.get("id") or "").strip()
            text = str(item.get("text") or "").strip()
            if not item_id or not text or item_id in seen:
                continue
            seen.add(item_id)
            rows.append(
                {
                    "content": text,
                    "custom_id": f"hindsight:{args.bank_id}:{item_id}",
                    "metadata": {
                        "source": "hindsight",
                        "original_id": item_id,
                        "bank_id": args.bank_id,
                        "date": item.get("date"),
                        "mentioned_at": item.get("mentioned_at"),
                        "occurred_start": item.get("occurred_start"),
                        "occurred_end": item.get("occurred_end"),
                        "confidence": None,
                        "status": item.get("state") or "valid",
                        "fact_type": item.get("fact_type"),
                        "proof_count": item.get("proof_count"),
                        "tags": item.get("tags") or [],
                        "imported_at": imported_at,
                    },
                }
            )
            if len(rows) >= wanted:
                break
        if len(rows) >= wanted:
            break

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("x", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    output.chmod(0o600)
    print(json.dumps({"sampled": len(rows), "source_total": total, "output": str(output)}))


def ingest_jsonl(args: argparse.Namespace) -> None:
    token = require_token(args.api_key_env)
    base = args.supermemory_url.rstrip("/")
    succeeded = 0
    failed = 0
    for line_number, raw in enumerate(Path(args.input).read_text(encoding="utf-8").splitlines(), 1):
        if not raw.strip():
            continue
        row = json.loads(raw)
        payload = {
            "content": str(row["content"]),
            "containerTag": args.container_tag,
            "customId": str(row.get("custom_id") or f"poc-line:{line_number}"),
            "metadata": row.get("metadata") or {},
        }
        try:
            request_json(
                "POST",
                f"{base}/v3/documents",
                token=token,
                payload=payload,
                timeout=args.timeout,
            )
            succeeded += 1
        except RuntimeError as exc:
            failed += 1
            print(f"line {line_number}: {exc}", file=sys.stderr)
            if not args.keep_going:
                raise
        time.sleep(args.delay)
    print(json.dumps({"accepted": succeeded, "failed": failed, "container_tag": args.container_tag}))


def ingest_files(args: argparse.Namespace) -> None:
    token = require_token(args.api_key_env)
    base = args.supermemory_url.rstrip("/")
    accepted = 0
    max_bytes = args.max_file_mib * 1024 * 1024
    for line_number, raw in enumerate(Path(args.manifest).read_text(encoding="utf-8").splitlines(), 1):
        if not raw.strip():
            continue
        row = json.loads(raw)
        path = Path(str(row["path"])).expanduser().resolve()
        if not path.is_file():
            raise RuntimeError(f"manifest line {line_number}: file does not exist: {path}")
        size = path.stat().st_size
        if size > max_bytes:
            raise RuntimeError(
                f"manifest line {line_number}: {path} is {size} bytes, above the {args.max_file_mib} MiB cap"
            )
        source = str(row.get("source") or "").strip().lower()
        if source not in {"devonthink", "obsidian"}:
            raise RuntimeError(f"manifest line {line_number}: source must be devonthink or obsidian")
        original_id = str(row.get("original_id") or "").strip()
        original_path = str(row.get("original_path") or path)
        stable_input = f"{source}\0{original_id}\0{original_path}".encode("utf-8")
        custom_id = str(row.get("custom_id") or f"{source}:{hashlib.sha256(stable_input).hexdigest()[:32]}")
        metadata = {
            "source": source,
            "original_id": original_id,
            "original_path": original_path,
            "date": str(row.get("date") or ""),
            "imported_at": datetime.now(timezone.utc).isoformat(),
            "confidence": str(row.get("confidence") or ""),
            "status": str(row.get("status") or "reviewed"),
        }
        response = request_multipart_json(
            f"{base}/v3/documents/file",
            token=token,
            fields={
                "containerTag": args.container_tag,
                "customId": custom_id,
                "filepath": original_path,
                "metadata": json.dumps(metadata, ensure_ascii=False),
            },
            file_path=path,
            timeout=args.timeout,
        )
        accepted += 1
        print(json.dumps({"line": line_number, "id": response.get("id"), "status": response.get("status")}))
        time.sleep(args.delay)
    print(json.dumps({"accepted": accepted, "container_tag": args.container_tag}))


def smoke(args: argparse.Namespace) -> None:
    token = require_token(args.api_key_env)
    base = args.supermemory_url.rstrip("/")
    marker = f"supermemory-poc-smoke-{int(time.time())}"
    container = args.container_tag
    add = request_json(
        "POST",
        f"{base}/v3/documents",
        token=token,
        payload={
            "content": f"{marker}: 中文跨文档检索烟测；authoritative source is the POC harness.",
            "containerTag": container,
            "customId": marker,
            "metadata": {"source": "poc-smoke", "status": "temporary"},
        },
        timeout=args.timeout,
    )
    document_id = str(add.get("id") or "")
    deadline = time.monotonic() + args.wait
    found = False
    while time.monotonic() < deadline:
        search = request_json(
            "POST",
            f"{base}/v4/search",
            token=token,
            payload={
                "q": marker,
                "containerTag": container,
                "searchMode": "documents",
                "limit": 5,
            },
            timeout=args.timeout,
        )
        found = bool(search.get("results"))
        if found:
            break
        time.sleep(2)
    profile = request_json(
        "POST",
        f"{base}/v4/profile",
        token=token,
        payload={"containerTag": container, "q": marker},
        timeout=args.timeout,
    )
    request_json(
        "POST",
        f"{base}/v4/conversations",
        token=token,
        payload={
            "conversationId": marker,
            "containerTags": [container],
            "messages": [
                {"role": "user", "content": "记住本轮只是隔离 POC 烟测。"},
                {"role": "assistant", "content": "已记录为临时测试，不作为正式事实。"},
            ],
            "metadata": {"source": "poc-smoke", "status": "temporary"},
        },
        timeout=args.timeout,
    )
    if not found:
        raise RuntimeError(f"document {document_id or marker} was not searchable within {args.wait}s")
    profile_shape_ok = isinstance(profile, dict) and "profile" in profile
    if not profile_shape_ok:
        raise RuntimeError("v4 profile response is missing the profile field expected by Hermes")
    print(
        json.dumps(
            {
                "document_id": document_id,
                "documents_search": "ok",
                "profile": "ok",
                "conversation_ingest": "accepted",
                "container_tag": container,
            }
        )
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)

    sample = commands.add_parser("sample-hindsight", help="Read a stratified observation sample")
    sample.add_argument("--hindsight-url", default="http://192.168.50.100:8888")
    sample.add_argument("--bank-id", default="personal-context")
    sample.add_argument("--limit", type=int, default=60)
    sample.add_argument("--page-size", type=int, default=10)
    sample.add_argument("--timeout", type=float, default=30.0)
    sample.add_argument("--output", required=True)
    sample.set_defaults(func=sample_hindsight)

    ingest = commands.add_parser("ingest-jsonl", help="Serially ingest reviewed JSONL rows")
    ingest.add_argument("--supermemory-url", required=True)
    ingest.add_argument("--container-tag", default="hermes_supermemory_lab")
    ingest.add_argument("--api-key-env", default="SUPERMEMORY_API_KEY")
    ingest.add_argument("--input", required=True)
    ingest.add_argument("--delay", type=float, default=0.5)
    ingest.add_argument("--timeout", type=float, default=60.0)
    ingest.add_argument("--keep-going", action="store_true")
    ingest.set_defaults(func=ingest_jsonl)

    files = commands.add_parser("ingest-files", help="Upload a reviewed DEVONthink/Obsidian manifest")
    files.add_argument("--supermemory-url", required=True)
    files.add_argument("--container-tag", default="hermes_supermemory_lab")
    files.add_argument("--api-key-env", default="SUPERMEMORY_API_KEY")
    files.add_argument("--manifest", required=True)
    files.add_argument("--max-file-mib", type=int, default=50)
    files.add_argument("--delay", type=float, default=1.0)
    files.add_argument("--timeout", type=float, default=120.0)
    files.set_defaults(func=ingest_files)

    check = commands.add_parser("smoke", help="Exercise the API paths used by Hermes")
    check.add_argument("--supermemory-url", required=True)
    check.add_argument("--container-tag", default="hermes_supermemory_lab")
    check.add_argument("--api-key-env", default="SUPERMEMORY_API_KEY")
    check.add_argument("--wait", type=float, default=120.0)
    check.add_argument("--timeout", type=float, default=60.0)
    check.set_defaults(func=smoke)
    return parser


def main() -> int:
    try:
        args = build_parser().parse_args()
        args.func(args)
        return 0
    except (KeyError, ValueError, OSError, RuntimeError, json.JSONDecodeError) as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


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