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

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


MESSAGE_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"


def make_ulid(index: int) -> str:
    alphabet = "0123456789ABCDEFGHJKMNPQRSTVWXYZ"
    value = index
    output = ["0"] * 26
    for position in range(25, -1, -1):
        output[position] = alphabet[value & 31]
        value >>= 5
    return "".join(output)


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)

    truncated = root / "truncated.jsonl"
    write_jsonl(truncated, record())
    truncated.write_bytes(truncated.read_bytes()[:-8])
    run(server, root / "truncated-data", "--import", str(truncated), success=False)

    crlf = root / "crlf.jsonl"
    crlf.write_bytes(source.read_bytes().replace(b"\n", b"\r\n", 1).replace(b"\r\n", b"\r\n\r\n", 1))
    run(server, root / "crlf-data", "--import", str(crlf), success=True)

    dangling = root / "dangling.jsonl"
    dangling_record = record()
    dangling_record["metadata"] = dict(dangling_record["metadata"])  # type: ignore[arg-type]
    dangling_record["metadata"]["relations"] = [  # type: ignore[index]
        {"type": "derived_from", "message_id": "01ARZ3NDEKTSV4RRFFQ69G5FAW"}
    ]
    write_jsonl(dangling, dangling_record)
    run(server, root / "dangling-data", "--import", str(dangling), success=False)

    duplicate = root / "duplicate.jsonl"
    duplicate.write_text(
        json.dumps({"type": "tiny-convmem-manifest", "format_version": 2})
        + "\n"
        + json.dumps(record("first"), separators=(",", ":"))
        + "\n"
        + json.dumps(record("last"), separators=(",", ":"))
        + "\n",
        encoding="utf-8",
    )
    duplicate_data = root / "duplicate-data"
    run(server, duplicate_data, "--import", str(duplicate), success=True)
    duplicate_export = root / "duplicate-export.jsonl"
    run(server, duplicate_data, "--export", str(duplicate_export), success=True)
    assert json.loads(duplicate_export.read_text(encoding="utf-8").splitlines()[1])["content"] == "last"

    over_count = root / "over-count.jsonl"
    with over_count.open("w", encoding="utf-8", newline="\n") as file:
        file.write(json.dumps({"type": "tiny-convmem-manifest", "format_version": 2}) + "\n")
        for index in range(100_001):
            file.write(json.dumps({"type": "DEL", "message_id": make_ulid(index)}) + "\n")
    run(server, root / "over-count-data", "--import", str(over_count), success=False)
    assert not (root / "over-count-data" / "active.import.tmp").exists()

    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()