"""Build the encrypted history snapshot served at /export/v1/messages.

Shape of a build
----------------
1. One read-only REPEATABLE READ transaction over both message tables, so the
   rolling window and the rows that follow come from the same instant.
2. The window: whole sessions, newest activity first, until the raw text
   reaches `navi_export_max_bytes`. Cutting on a session boundary keeps a
   conversation from arriving half a story.
3. Rows are streamed through a cursor and serialised one JSON object per line.
   The plaintext gzip is built deterministically (fixed mtime, sorted keys) so
   the same input always yields the same `content_id` — that is what lets the
   consumer answer "nothing changed" without downloading anything.
4. The plaintext is encrypted to `navi_export_age_recipient` and the plaintext is
   deleted. Only ciphertext stays on disk.

Nothing here reads `users`, `mcp_user_keys`, `api_tokens` or attachments.
Adding a field is a schema change: bump EXPORT_VERSION and say so in the
manifest, because the consumer validates `fields`.
"""

from __future__ import annotations

import gzip
import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path

import structlog

from navi.config import settings

log = structlog.get_logger()

EXPORT_VERSION = 1
FIELDS = ["v", "session", "seq", "role", "content", "ts"]
MANIFEST_NAME = "manifest.json"
FILE_NAME = "messages-v1.jsonl.gz.age"

# Whole sessions, newest first, while the cumulative text stays under budget.
_WINDOW_SQL = """
WITH msgs AS (
    SELECT session_id, role, content, created_at, is_summary, is_compression
    FROM session_messages
    UNION ALL
    SELECT session_id, role, content, created_at, is_summary, is_compression
    FROM session_messages_archive
),
buckets AS (
    SELECT session_id,
           max(created_at) AS last_ts,
           sum(octet_length(content))::bigint AS text_bytes
    FROM msgs
    WHERE role = ANY($1::text[])
      AND content IS NOT NULL AND content <> ''
      AND NOT coalesce(is_summary, false)
      AND NOT coalesce(is_compression, false)
    GROUP BY session_id
),
ordered AS (
    SELECT session_id, last_ts, text_bytes,
           sum(text_bytes) OVER (ORDER BY last_ts DESC NULLS LAST, session_id) AS cum
    FROM buckets
)
SELECT session_id
FROM ordered
WHERE cum - text_bytes < $2
ORDER BY last_ts DESC NULLS LAST, session_id
"""

_SESSIONS_SQL = """
SELECT count(*) FROM (
    SELECT session_id
    FROM (
        SELECT session_id, role, content, created_at, is_summary, is_compression
        FROM session_messages
        UNION ALL
        SELECT session_id, role, content, created_at, is_summary, is_compression
        FROM session_messages_archive
    ) m
    WHERE role = ANY($1::text[])
      AND content IS NOT NULL AND content <> ''
      AND NOT coalesce(is_summary, false)
      AND NOT coalesce(is_compression, false)
    GROUP BY session_id
) t
"""

_ROWS_SQL = """
SELECT session_id, sequence_number, role, content, created_at
FROM (
    SELECT session_id, sequence_number, role, content, created_at, is_summary, is_compression
    FROM session_messages
    UNION
    SELECT session_id, sequence_number, role, content, created_at, is_summary, is_compression
    FROM session_messages_archive
) m
WHERE session_id = ANY($1::text[])
  AND role = ANY($2::text[])
  AND content IS NOT NULL AND content <> ''
  AND NOT coalesce(is_summary, false)
  AND NOT coalesce(is_compression, false)
ORDER BY created_at NULLS LAST, session_id, sequence_number
"""


def export_paths(out_dir: Path | None = None) -> tuple[Path, Path]:
    """(age file, manifest file) in the export directory."""
    base = Path(out_dir if out_dir is not None else settings.navi_export_dir)
    return base / FILE_NAME, base / MANIFEST_NAME


def read_manifest(out_dir: Path | None = None) -> dict | None:
    """The manifest on disk, or None when no export has been built yet."""
    _, manifest_path = export_paths(out_dir)
    try:
        return json.loads(manifest_path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None


def _record(row) -> dict:
    created = row["created_at"]
    ts = created.astimezone(timezone.utc).isoformat() if created else None
    return {
        "v": EXPORT_VERSION,
        "session": row["session_id"],
        "seq": row["sequence_number"],
        "role": row["role"],
        "content": row["content"],
        "ts": ts,
    }


async def build_export(pool, out_dir: Path | None = None) -> dict:
    """Build (or confirm) the snapshot and return the manifest.

    Raises when the export is not configured or the recipient key is malformed;
    the caller is a background loop that logs and retries on the next tick.
    """
    import pyrage

    if not settings.navi_export_configured:
        raise RuntimeError("export is not configured (NAVI_EXPORT_AGE_RECIPIENT / NAVI_EXPORT_TOKEN)")

    recipient = pyrage.x25519.Recipient.from_str(settings.navi_export_age_recipient.strip())
    roles = settings.navi_export_roles_list
    # 0 (or negative) disables the cap: the window then covers all history.
    budget = settings.navi_export_max_bytes if settings.navi_export_max_bytes > 0 else 2**62

    age_path, manifest_path = export_paths(out_dir)
    base = age_path.parent
    base.mkdir(parents=True, exist_ok=True)

    plain_tmp = base / (FILE_NAME + ".plain.tmp")
    age_tmp = base / (FILE_NAME + ".tmp")

    now = datetime.now(timezone.utc)
    rows = 0
    text_bytes = 0
    first_ts: datetime | None = None
    last_ts: datetime | None = None
    session_ids: list[str] = []
    total_sessions = 0

    digest = hashlib.sha256()
    try:
        async with pool.acquire() as conn, conn.transaction(
            isolation="repeatable_read", readonly=True
        ):
            session_ids = [
                r["session_id"] for r in await conn.fetch(_WINDOW_SQL, roles, budget)
            ]
            total_sessions = await conn.fetchval(_SESSIONS_SQL, roles) or 0

            # Blocking file I/O in a background job, once every few hours:
            # an executor hop here would buy nothing.
            with open(plain_tmp, "wb") as plain_fh, gzip.GzipFile(  # noqa: ASYNC230
                filename="",
                mode="wb",
                fileobj=plain_fh,
                compresslevel=9,
                mtime=0,
            ) as gz:
                if session_ids:
                    async for row in conn.cursor(_ROWS_SQL, session_ids, roles, prefetch=500):
                        line = json.dumps(
                            _record(row), ensure_ascii=False, sort_keys=True, separators=(",", ":")
                        ) + "\n"
                        chunk = line.encode("utf-8")
                        gz.write(chunk)
                        digest.update(chunk)
                        rows += 1
                        text_bytes += len(row["content"].encode("utf-8"))
                        created = row["created_at"]
                        if created is not None:
                            if first_ts is None or created < first_ts:
                                first_ts = created
                            if last_ts is None or created > last_ts:
                                last_ts = created
    except BaseException:
        plain_tmp.unlink(missing_ok=True)
        raise

    content_id = digest.hexdigest()
    previous = read_manifest(out_dir)

    # Same corpus as last time: keep the ciphertext (and its ETag) untouched.
    # Only the manifest's built_at moves, which is what the consumer ignores.
    if (
        previous
        and previous.get("content_id") == content_id
        and age_path.exists()
        and previous.get("file", {}).get("name") == FILE_NAME
    ):
        plain_tmp.unlink(missing_ok=True)
        manifest = {**previous, "built_at": now.isoformat()}
        _write_manifest(manifest_path, manifest)
        log.info("export.unchanged", content_id=content_id[:12], rows=rows)
        return manifest

    try:
        with open(plain_tmp, "rb") as src, open(age_tmp, "wb") as dst:  # noqa: ASYNC230
            pyrage.encrypt_io(src, dst, [recipient])
        file_bytes = age_tmp.stat().st_size
        file_sha = _sha256_file(age_tmp)
        os.replace(age_tmp, age_path)
    finally:
        plain_tmp.unlink(missing_ok=True)
        age_tmp.unlink(missing_ok=True)

    manifest = {
        "export_version": EXPORT_VERSION,
        "built_at": now.isoformat(),
        "content_id": content_id,
        "file": {"name": FILE_NAME, "bytes": file_bytes, "sha256": file_sha},
        "rows": rows,
        "sessions": len(session_ids),
        "sessions_total": total_sessions,
        "text_bytes": text_bytes,
        "fields": FIELDS,
        "roles": roles,
        "window": {
            "from": first_ts.astimezone(timezone.utc).isoformat() if first_ts else None,
            "to": last_ts.astimezone(timezone.utc).isoformat() if last_ts else None,
            "truncated": len(session_ids) < total_sessions,
            "max_bytes": budget,
        },
        "age_recipient": settings.navi_export_age_recipient.strip(),
    }
    _write_manifest(manifest_path, manifest)
    log.info(
        "export.built",
        rows=rows,
        sessions=len(session_ids),
        text_bytes=text_bytes,
        file_bytes=file_bytes,
        content_id=content_id[:12],
    )
    return manifest


def _write_manifest(path: Path, manifest: dict) -> None:
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    os.replace(tmp, path)


def _sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()