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

import argparse
import json
import subprocess
import tempfile
from pathlib import Path


MESSAGE_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"


def record(content: str = "valid", **changes: object) -> dict[str, object]:
    value: dict[str, object] = {
        "message_id": MESSAGE_ID,
        "session_id": "escape-session",
        "timestamp": 0,
        "updated_at": 9223372036854775807,
        "revision": 2147483647,
        "role": "user",
        "content": content,
        "metadata": {
            "tool_call_id": "tool\\\"\n",
            "tags": ["tag\\\"\n"],
            "memory_tier": "short",
            "memory_kind": "conversation",
            "importance": 100,
            "summary": "summary\\\"\n\u3042",
            "relations": [],
            "source_message_ids": [],
        },
    }
    value.update(changes)
    return value


def write_jsonl(path: Path, item: dict[str, object], final_newline: bool = True) -> None:
    text = json.dumps({"type": "tiny-convmem-manifest", "format_version": 2}) + "\n"
    text += json.dumps(item, ensure_ascii=True, separators=(",", ":"))
    path.write_text(text + ("\n" if final_newline else ""), encoding="utf-8")


def run(server: Path, data: Path, *arguments: str, success: bool) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(
        [str(server), "--data-dir", str(data), *arguments],
        capture_output=True,
        text=True,
    )
    assert (result.returncode == 0) is success, (
        result.returncode,
        result.args,
        result.stdout,
        result.stderr,
    )
    return result


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", required=True, type=Path)
    args = parser.parse_args()
    server = args.server.resolve()
    root = Path(tempfile.mkdtemp(prefix="tiny-convmem-import-e2e-"))

    prefix = "prefix\nquote=\" slash=\\ unicode=\u3042 "
    escaped = prefix + "x" * (1024 * 1024 - len(prefix.encode("utf-8")))
    assert len(escaped.encode("utf-8")) == 1024 * 1024
    source = root / "source.jsonl"
    write_jsonl(source, record(escaped))
    first = root / "first"
    run(server, first, "--import", str(source), success=True)
    exported = root / "export.jsonl"
    run(server, first, "--export", str(exported), success=True)
    second = root / "second"
    run(server, second, "--import", str(exported), success=True)
    exported_again = root / "export-again.jsonl"
    run(server, second, "--export", str(exported_again), success=True)
    assert exported.read_bytes() == exported_again.read_bytes()
    assert json.loads(exported.read_text(encoding="utf-8").splitlines()[1])["content"] == escaped

    invalid_values: list[tuple[str, object]] = [
        ("timestamp", -1),
        ("timestamp", 9223372036854775808),
        ("timestamp", 1.5),
        ("timestamp", 1e3),
        ("revision", 0),
        ("revision", 2147483648),
    ]
    for index, (field, value) in enumerate(invalid_values):
        invalid = root / f"invalid-{index}.jsonl"
        write_jsonl(invalid, record(**{field: value}))
        run(server, root / f"invalid-data-{index}", "--import", str(invalid), success=False)

    malformed = root / "embedded-nul.jsonl"
    write_jsonl(malformed, record("before\0after"))
    run(server, root / "nul-data", "--import", str(malformed), success=False)

    long_line = root / "long-line.jsonl"
    long_line.write_text(
        '{"type":"tiny-convmem-manifest","format_version":2}\n' + "x" * (1024 * 1024 + 4097),
        encoding="utf-8",
    )
    run(server, root / "long-data", "--import", str(long_line), success=False)

    atomic_target = root / "atomic.jsonl"
    atomic_target.write_text("old backup", encoding="utf-8")
    run(server, first, "--export", str(atomic_target), success=True)
    assert "tiny-convmem-manifest" in atomic_target.read_text(encoding="utf-8")
    assert not Path(str(atomic_target) + ".tmp").exists()

    protected = root / "protected.jsonl"
    protected.write_text("existing backup", encoding="utf-8")
    Path(str(protected) + ".tmp").mkdir()
    run(server, first, "--export", str(protected), success=False)
    assert protected.read_text(encoding="utf-8") == "existing backup"


if __name__ == "__main__":
    main()