// One root session's memory on disk: main/ and tree/ day logs (append-only), view.json, and a lock.
import {
    appendFile,
    mkdir,
    open,
    readFile,
    readdir,
    rename,
    rm,
    writeFile,
} from "node:fs/promises";
import { join } from "node:path";
import { z } from "zod";
import { bytes, key, pieces, type Kind, type Line, type View } from "./optchat-core.js";

const kind = z.enum(["user", "agent", "tool", "echo"]);
const record = z.object({
    i: z.number().int().nonnegative(),
    kind,
    text: z.string(),
    size: z.number(),
    date: z.string(),
    /** The history part this message came from; it is never logged twice. */
    src: z.string(),
});
const node = z.object({
    l: z.number().int().nonnegative(),
    i: z.number().int().nonnegative(),
    text: z.string(),
    size: z.number(),
});
const line = z.tuple([z.number().int().nonnegative(), z.number().int().nonnegative()]);
const view = z.object({ lines: z.array(line), batching: z.boolean() });
export type Message = z.infer<typeof record>;

async function lines(directory: string) {
    const names = await readdir(directory).catch(() => []);
    const out: string[] = [];
    for (const name of names.filter((name) => name.endsWith(".jsonl")).sort())
        out.push(...(await readFile(join(directory, name), "utf8")).split("\n"));
    // A torn last line from a crash is skipped; its message or node is redone.
    return out
        .filter((text) => text.trim())
        .flatMap((text) => {
            try {
                return [JSON.parse(text)];
            } catch {
                return [];
            }
        });
}

async function lock(file: string) {
    for (let attempt = 0; attempt < 2; attempt++) {
        try {
            const handle = await open(file, "wx");
            await handle.writeFile(String(process.pid));
            await handle.close();
            return;
        } catch (error) {
            if (!(error instanceof Error && "code" in error && error.code === "EEXIST"))
                throw error;
        }
        const pid = Number(await readFile(file, "utf8").catch(() => ""));
        if (Number.isInteger(pid) && pid > 0 && pid !== process.pid) {
            try {
                process.kill(pid, 0);
                throw new Error(`optchat memory is in use by process ${pid}`);
            } catch (error) {
                if (!(error instanceof Error && "code" in error && error.code === "ESRCH"))
                    throw error;
            }
        } else if (pid === process.pid) throw new Error("optchat memory is already open");
        await rm(file, { force: true });
    }
    throw new Error("optchat memory lock is contended");
}

export class Memory {
    readonly messages: Message[] = [];
    readonly #tree = new Map<string, string>();
    readonly #sources = new Set<string>();
    #writes: Promise<void> = Promise.resolve();
    view: View = { lines: [], batching: false };

    private constructor(readonly directory: string) {}

    /** Takes the lock and loads the log, tree and saved view. */
    static async open(directory: string) {
        await mkdir(join(directory, "main"), { recursive: true });
        await mkdir(join(directory, "tree"), { recursive: true });
        await lock(join(directory, "lock"));
        const memory = new Memory(directory);
        try {
            for (const value of await lines(join(directory, "main"))) {
                const parsed = record.safeParse(value);
                if (!parsed.success || parsed.data.i !== memory.messages.length) continue;
                memory.messages.push(parsed.data);
                memory.#sources.add(parsed.data.src);
            }
            for (const value of await lines(join(directory, "tree"))) {
                const parsed = node.safeParse(value);
                if (parsed.success && (parsed.data.l > 0 || parsed.data.i < memory.messages.length))
                    memory.#tree.set(key([parsed.data.l, parsed.data.i]), parsed.data.text);
            }
            const saved = await readFile(join(directory, "view.json"), "utf8").catch(
                () => undefined,
            );
            if (saved !== undefined) memory.view = view.parse(JSON.parse(saved));
            // Messages logged after the last saved view (a crash in between) are appended in order.
            const covered = memory.view.lines.reduce((sum, [l]) => sum + 2 ** l, 0);
            for (let i = covered; i < memory.messages.length; i++) memory.view.lines.push([0, i]);
        } catch (error) {
            await memory.close();
            throw error;
        }
        return memory;
    }

    has(source: string) {
        return this.#sources.has(source);
    }

    text(line: Line) {
        return this.#tree.get(key(line));
    }

    nodes() {
        return [...this.#tree.keys()].map((name): Line => {
            const [l, i] = name.split(":").map(Number);
            return [l!, i!];
        });
    }

    /** Logs one history part as one or more messages, in order; returns their ids. */
    log(kind: Kind, text: string, source: string) {
        if (this.#sources.has(source)) return [];
        this.#sources.add(source);
        const date = new Date().toISOString();
        return pieces(kind, text).map((text, k) => {
            const message: Message = {
                i: this.messages.length,
                kind,
                text,
                size: bytes(text),
                date,
                src: k === 0 ? source : `${source}#${k}`,
            };
            this.messages.push(message);
            this.#append("main", date, message);
            return message.i;
        });
    }

    node([l, i]: Line, text: string) {
        this.#tree.set(key([l, i]), text);
        this.#append("tree", new Date().toISOString(), { l, i, text, size: bytes(text) });
    }

    saveView() {
        const snapshot = JSON.stringify(this.view);
        this.#write(async () => {
            const file = join(this.directory, "view.json");
            await writeFile(`${file}.tmp`, snapshot);
            await rename(`${file}.tmp`, file);
        });
    }

    /** Resolves once every write so far reached the file. */
    flushed() {
        return this.#writes;
    }

    async close() {
        await this.#writes;
        await rm(join(this.directory, "lock"), { force: true });
    }

    #append(directory: "main" | "tree", date: string, value: Message | z.infer<typeof node>) {
        const file = join(this.directory, directory, `${date.slice(0, 10)}.jsonl`);
        const text = JSON.stringify(value) + "\n";
        this.#write(() => appendFile(file, text));
    }

    #write(write: () => Promise<void>) {
        // One writer, in order; a failed write is reported and later writes still run.
        this.#writes = this.#writes
            .then(write)
            .catch((error) => console.error("optchat write failed", error));
    }
}