import { createHash, randomUUID } from "node:crypto";
import fs from "node:fs";
import path from "node:path";
import { CURRENT_SESSION_VERSION } from "@mariozechner/pi-coding-agent";

export type SessionRecord = {
  sessionKey: string;
  keyHash: string;
  sessionId: string;
  sessionFile: string;
  updatedAt: number;
};

type IndexFileShapeV1 = {
  version: 1;
  updatedAt: number;
  sessions: Record<
    string,
    {
      keyHash: string;
      sessionId: string;
      sessionFile: string;
      updatedAt: number;
    }
  >;
};

function sha256Hex(input: string): string {
  return createHash("sha256").update(input).digest("hex");
}

function safeRename(from: string, to: string): void {
  try {
    fs.renameSync(from, to);
  } catch {
    // best-effort
  }
}

function ensureDir(dirPath: string): void {
  fs.mkdirSync(dirPath, { recursive: true });
}

function readJsonFile(filePath: string): unknown {
  const raw = fs.readFileSync(filePath, "utf-8");
  return JSON.parse(raw) as unknown;
}

function writeFileAtomic(filePath: string, content: string, mode: number): void {
  const dir = path.dirname(filePath);
  ensureDir(dir);
  const tmpPath = `${filePath}.tmp`;
  fs.writeFileSync(tmpPath, content, { encoding: "utf-8", mode });
  fs.renameSync(tmpPath, filePath);
}

function ensureTranscriptHeader(params: { sessionFile: string; sessionId: string }): void {
  if (fs.existsSync(params.sessionFile)) {
    return;
  }
  ensureDir(path.dirname(params.sessionFile));
  const header = {
    type: "session",
    version: CURRENT_SESSION_VERSION,
    id: params.sessionId,
    timestamp: new Date().toISOString(),
    cwd: process.cwd(),
  };
  fs.writeFileSync(params.sessionFile, `${JSON.stringify(header)}\n`, {
    encoding: "utf-8",
    mode: 0o600,
  });
}

function decodeIndexFileShapeV1(value: unknown): IndexFileShapeV1 | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const rec = value as Record<string, unknown>;
  if (rec.version !== 1) return null;
  if (!("sessions" in rec) || !rec.sessions || typeof rec.sessions !== "object" || Array.isArray(rec.sessions)) {
    return null;
  }
  const updatedAt = typeof rec.updatedAt === "number" ? rec.updatedAt : Date.now();
  const sessions = rec.sessions as Record<string, unknown>;
  const out: IndexFileShapeV1["sessions"] = {};
  for (const [sessionKey, entry] of Object.entries(sessions)) {
    if (!entry || typeof entry !== "object" || Array.isArray(entry)) continue;
    const e = entry as Record<string, unknown>;
    const keyHash = typeof e.keyHash === "string" ? e.keyHash : "";
    const sessionId = typeof e.sessionId === "string" ? e.sessionId : "";
    const sessionFile = typeof e.sessionFile === "string" ? e.sessionFile : "";
    const entryUpdatedAt = typeof e.updatedAt === "number" ? e.updatedAt : updatedAt;
    if (!sessionKey || !keyHash || !sessionId || !sessionFile) continue;
    out[sessionKey] = { keyHash, sessionId, sessionFile, updatedAt: entryUpdatedAt };
  }
  return { version: 1, updatedAt, sessions: out };
}

export class SessionStore {
  private readonly sessionsDir: string;
  private readonly indexPath: string;
  private readonly transcriptsDir: string;
  private readonly archiveDir: string;
  private readonly cache = new Map<string, SessionRecord>();

  constructor(private readonly agentDir: string) {
    this.sessionsDir = path.join(agentDir, ".pi-bot", "sessions");
    this.indexPath = path.join(this.sessionsDir, "index.json");
    this.transcriptsDir = path.join(this.sessionsDir, "transcripts");
    this.archiveDir = path.join(this.sessionsDir, "archive");
  }

  load(): void {
    this.cache.clear();

    if (!fs.existsSync(this.indexPath)) {
      return;
    }

    try {
      const decoded = decodeIndexFileShapeV1(readJsonFile(this.indexPath));
      if (!decoded) {
        throw new Error("invalid index.json schema");
      }
      for (const [sessionKey, entry] of Object.entries(decoded.sessions)) {
        this.cache.set(sessionKey, {
          sessionKey,
          keyHash: entry.keyHash,
          sessionId: entry.sessionId,
          sessionFile: entry.sessionFile,
          updatedAt: entry.updatedAt,
        });
      }
    } catch {
      // Keep a corrupt copy for debugging instead of crashing startup.
      safeRename(this.indexPath, `${this.indexPath}.corrupt-${Date.now()}`);
      this.cache.clear();
    }
  }

  listSessionKeys(): string[] {
    return Array.from(this.cache.keys());
  }

  get(sessionKey: string): SessionRecord | undefined {
    return this.cache.get(sessionKey);
  }

  ensureSession(sessionKey: string): SessionRecord {
    const existing = this.cache.get(sessionKey);
    if (existing) {
      // If the transcript was deleted externally, recreate the header so the
      // session remains recoverable.
      ensureTranscriptHeader({ sessionFile: existing.sessionFile, sessionId: existing.sessionId });
      return existing;
    }

    const now = Date.now();
    const keyHash = sha256Hex(sessionKey);
    const sessionId = randomUUID();
    const sessionFile = path.join(this.transcriptsDir, `${keyHash}-${sessionId}.jsonl`);
    const record: SessionRecord = {
      sessionKey,
      keyHash,
      sessionId,
      sessionFile,
      updatedAt: now,
    };
    this.cache.set(sessionKey, record);
    ensureTranscriptHeader({ sessionFile, sessionId });
    this.saveIndex();
    return record;
  }

  resetSession(sessionKey: string): SessionRecord {
    const prior = this.cache.get(sessionKey);
    if (prior && fs.existsSync(prior.sessionFile)) {
      ensureDir(this.archiveDir);
      const baseName = path.basename(prior.sessionFile);
      const archivedBase = baseName.endsWith(".jsonl") ? baseName : `${baseName}.jsonl`;
      const archivedPathCandidate = path.join(this.archiveDir, archivedBase);
      const archivedPath = fs.existsSync(archivedPathCandidate)
        ? path.join(this.archiveDir, `${prior.keyHash}-${prior.sessionId}-${Date.now()}.jsonl`)
        : archivedPathCandidate;
      safeRename(prior.sessionFile, archivedPath);
    }

    // Keep keyHash stable for this sessionKey, rotate sessionId + sessionFile.
    const now = Date.now();
    const keyHash = prior?.keyHash ?? sha256Hex(sessionKey);
    const sessionId = randomUUID();
    const sessionFile = path.join(this.transcriptsDir, `${keyHash}-${sessionId}.jsonl`);
    const next: SessionRecord = {
      sessionKey,
      keyHash,
      sessionId,
      sessionFile,
      updatedAt: now,
    };
    this.cache.set(sessionKey, next);
    ensureTranscriptHeader({ sessionFile, sessionId });
    this.saveIndex();
    return next;
  }

  private saveIndex(): void {
    const now = Date.now();
    const sessions: IndexFileShapeV1["sessions"] = {};
    for (const [sessionKey, entry] of this.cache.entries()) {
      sessions[sessionKey] = {
        keyHash: entry.keyHash,
        sessionId: entry.sessionId,
        sessionFile: entry.sessionFile,
        updatedAt: entry.updatedAt,
      };
    }
    const payload: IndexFileShapeV1 = {
      version: 1,
      updatedAt: now,
      sessions,
    };
    writeFileAtomic(this.indexPath, JSON.stringify(payload, null, 2), 0o600);
  }
}