/**
 * Chat session storage — sqlite-backed. One row per session, one row
 * per message; message.blocks is JSON-encoded.
 *
 * Lives in <data>/workflow.db alongside the rest of Forge's state.
 */

import { randomUUID } from 'node:crypto';
import { join } from 'node:path';
import { getDb } from '../../src/core/db/database';
import { getDataDir } from '../dirs';
import type { ContentBlock, Message, Role, Session } from './types';

function db() {
  return getDb(join(getDataDir(), 'workflow.db'));
}

interface SessionRow {
  id: string;
  title: string | null;
  created_at: number;
  updated_at: number;
  model: string | null;
  provider: string | null;
  system_prompt: string | null;
  meta: string | null;
}

interface MessageRow {
  id: string;
  session_id: string;
  role: Role;
  blocks: string;
  ts: number;
  error: string | null;
}

let initialized = false;

function ensureSchema(): void {
  if (initialized) return;
  const conn = db();
  conn.exec(`
    CREATE TABLE IF NOT EXISTS chat_sessions (
      id TEXT PRIMARY KEY,
      title TEXT,
      created_at INTEGER NOT NULL,
      updated_at INTEGER NOT NULL,
      model TEXT,
      provider TEXT,
      system_prompt TEXT,
      meta TEXT
    );
    CREATE INDEX IF NOT EXISTS idx_chat_sessions_updated ON chat_sessions(updated_at DESC);

    CREATE TABLE IF NOT EXISTS chat_messages (
      id TEXT PRIMARY KEY,
      session_id TEXT NOT NULL,
      role TEXT NOT NULL,
      blocks TEXT NOT NULL,
      ts INTEGER NOT NULL,
      error TEXT,
      FOREIGN KEY(session_id) REFERENCES chat_sessions(id) ON DELETE CASCADE
    );
    CREATE INDEX IF NOT EXISTS idx_chat_messages_session ON chat_messages(session_id, ts);
  `);
  initialized = true;
}

function rowToSession(r: SessionRow): Session {
  return {
    id: r.id,
    title: r.title,
    created_at: r.created_at,
    updated_at: r.updated_at,
    model: r.model,
    provider: r.provider,
    system_prompt: r.system_prompt,
    meta: r.meta ? safeParse(r.meta) : undefined,
  };
}

function rowToMessage(r: MessageRow): Message {
  return {
    id: r.id,
    session_id: r.session_id,
    role: r.role,
    blocks: safeParse<ContentBlock[]>(r.blocks) ?? [],
    ts: r.ts,
    error: r.error ?? undefined,
  };
}

function safeParse<T = unknown>(s: string): T | undefined {
  try { return JSON.parse(s) as T; } catch { return undefined; }
}

// ─── Sessions ────────────────────────────────────────────

export function createSession(opts: {
  title?: string;
  model?: string;
  provider?: string;
  system_prompt?: string;
  meta?: Record<string, unknown>;
}): Session {
  ensureSchema();
  const now = Date.now();
  const row: SessionRow = {
    id: randomUUID(),
    title: opts.title ?? null,
    created_at: now,
    updated_at: now,
    model: opts.model ?? null,
    provider: opts.provider ?? null,
    system_prompt: opts.system_prompt ?? null,
    meta: opts.meta ? JSON.stringify(opts.meta) : null,
  };
  db().prepare(`
    INSERT INTO chat_sessions (id, title, created_at, updated_at, model, provider, system_prompt, meta)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
  `).run(row.id, row.title, row.created_at, row.updated_at, row.model, row.provider, row.system_prompt, row.meta);
  return rowToSession(row);
}

export function getSession(id: string): Session | null {
  ensureSchema();
  const r = db().prepare(`SELECT * FROM chat_sessions WHERE id = ?`).get(id) as SessionRow | undefined;
  return r ? rowToSession(r) : null;
}

export function listSessions(limit = 50): Session[] {
  ensureSchema();
  const rows = db().prepare(`
    SELECT * FROM chat_sessions ORDER BY updated_at DESC LIMIT ?
  `).all(limit) as SessionRow[];
  return rows.map(rowToSession);
}

export function updateSession(id: string, patch: Partial<Pick<Session, 'title' | 'model' | 'provider' | 'system_prompt' | 'meta'>>): boolean {
  ensureSchema();
  const cur = getSession(id);
  if (!cur) return false;
  const merged = {
    title: patch.title ?? cur.title,
    model: patch.model ?? cur.model,
    provider: patch.provider ?? cur.provider,
    system_prompt: patch.system_prompt ?? cur.system_prompt,
    meta: patch.meta !== undefined ? patch.meta : cur.meta,
  };
  db().prepare(`
    UPDATE chat_sessions SET title=?, model=?, provider=?, system_prompt=?, meta=?, updated_at=?
    WHERE id=?
  `).run(
    merged.title, merged.model, merged.provider, merged.system_prompt,
    merged.meta ? JSON.stringify(merged.meta) : null,
    Date.now(), id,
  );
  return true;
}

export function deleteSession(id: string): boolean {
  ensureSchema();
  const r = db().prepare(`DELETE FROM chat_sessions WHERE id = ?`).run(id);
  return r.changes > 0;
}

/**
 * Empty all messages in a session WITHOUT deleting the session row.
 * Use for the main chat's "clear" — keep the session, drop history.
 */
export function clearSessionMessages(id: string): number {
  ensureSchema();
  const r = db().prepare(`DELETE FROM chat_messages WHERE session_id = ?`).run(id);
  touchSession(id);
  return r.changes;
}

/**
 * Find-or-create the user's persistent main session. Identified by
 * meta.kind === 'main'. Only one is ever returned (oldest wins if
 * multiple — caller can collapse later). Job dispatches and the "new
 * conversation" button create sessions with meta.kind = 'temp', so
 * the main session is naturally pinned and never auto-deleted.
 */
export function ensureMainSession(): Session {
  ensureSchema();
  const existing = db().prepare(
    `SELECT * FROM chat_sessions WHERE json_extract(meta, '$.kind') = 'main' ORDER BY created_at ASC LIMIT 1`,
  ).get() as SessionRow | undefined;
  if (existing) return rowToSession(existing);
  return createSession({ title: 'Main conversation', meta: { kind: 'main' } });
}

/**
 * Fork a session: copy its current messages into a brand-new session.
 * The fork starts as a temp session so it can be deleted later. Optional
 * upToTs lets the caller branch from a specific point (everything with
 * ts <= upToTs is copied).
 */
export function forkSession(srcId: string, opts: { title?: string; upToTs?: number } = {}): Session | null {
  ensureSchema();
  const src = getSession(srcId);
  if (!src) return null;
  const title = opts.title || `${src.title || 'Conversation'} (fork)`;
  const fork = createSession({
    title,
    model: src.model || undefined,
    provider: src.provider || undefined,
    system_prompt: src.system_prompt || undefined,
    meta: { kind: 'temp', forked_from: srcId },
  });
  const ts = opts.upToTs ?? Date.now();
  const msgs = db().prepare(
    `SELECT * FROM chat_messages WHERE session_id = ? AND ts <= ? ORDER BY ts ASC`,
  ).all(srcId, ts) as MessageRow[];
  const ins = db().prepare(`
    INSERT INTO chat_messages (id, session_id, role, blocks, ts, error)
    VALUES (?, ?, ?, ?, ?, ?)
  `);
  for (const m of msgs) {
    ins.run(randomUUID(), fork.id, m.role, m.blocks, m.ts, m.error);
  }
  touchSession(fork.id);
  return fork;
}

function touchSession(id: string): void {
  db().prepare(`UPDATE chat_sessions SET updated_at = ? WHERE id = ?`).run(Date.now(), id);
}

// ─── Messages ────────────────────────────────────────────

// Cap any single field to 64KB before it hits sqlite. LLM proxies (LiteLLM,
// vLLM) sometimes return error payloads listing every message that failed
// validation — a 188-turn chat trips 6000+ entries and the JSON balloons
// to 25+ MB. One bad row drags out every UI page-load that touches it,
// so truncate at write time.
const FIELD_BYTE_CAP = 64 * 1024;
function capLargeText(s: string | undefined | null): string | null {
  if (s == null) return null;
  if (s.length <= FIELD_BYTE_CAP) return s;
  return s.slice(0, FIELD_BYTE_CAP) + `\n… (truncated, full length ${s.length} chars)`;
}

export function appendMessage(opts: {
  session_id: string;
  role: Role;
  blocks: ContentBlock[];
  error?: string;
}): Message {
  ensureSchema();
  // Truncate any text block that's gone runaway (e.g. a streamed error
  // message stuffed into the assistant turn for UI display).
  const cappedBlocks = opts.blocks.map(b => {
    if (b.type === 'text' && b.text.length > FIELD_BYTE_CAP) {
      return { ...b, text: capLargeText(b.text)! };
    }
    return b;
  });
  const row: MessageRow = {
    id: randomUUID(),
    session_id: opts.session_id,
    role: opts.role,
    blocks: JSON.stringify(cappedBlocks),
    ts: Date.now(),
    error: capLargeText(opts.error),
  };
  db().prepare(`
    INSERT INTO chat_messages (id, session_id, role, blocks, ts, error)
    VALUES (?, ?, ?, ?, ?, ?)
  `).run(row.id, row.session_id, row.role, row.blocks, row.ts, row.error);
  touchSession(opts.session_id);
  return rowToMessage(row);
}

export function listMessages(session_id: string, opts?: { limit?: number; after_ts?: number }): Message[] {
  ensureSchema();
  const limit = opts?.limit ?? 500;
  const after = opts?.after_ts ?? 0;
  // Take the LAST N matching (DESC + LIMIT) then reverse to chronological
  // order for the caller. ASC + LIMIT was returning the OLDEST N, so a
  // session with >limit messages looked frozen at message #N's timestamp.
  const rows = db().prepare(`
    SELECT * FROM chat_messages WHERE session_id = ? AND ts > ?
    ORDER BY ts DESC LIMIT ?
  `).all(session_id, after, limit) as MessageRow[];
  return rows.map(rowToMessage).reverse();
}

/** Last N messages in chronological order — used by agent-loop to cap LLM context. */
export function listRecentMessages(session_id: string, limit: number): Message[] {
  ensureSchema();
  const rows = db().prepare(`
    SELECT * FROM chat_messages WHERE session_id = ?
    ORDER BY ts DESC LIMIT ?
  `).all(session_id, limit) as MessageRow[];
  return rows.map(rowToMessage).reverse();
}

/**
 * Take the most recent messages, stopping when either the message-count
 * budget OR the token-estimate budget would be exceeded. Walks
 * newest-first so the most recent dialogue is always kept; returns
 * chronological order for the LLM history slot.
 *
 * msgBudget is enforced via SQL LIMIT (cheap). tokenBudget is enforced
 * via the caller-supplied estimator (decoupled to avoid pulling the
 * token-estimator into the storage layer).
 */
export function listMessagesCapped(
  session_id: string,
  msgBudget: number,
  tokenBudget: number,
  estimateTokens: (m: Message) => number,
): Message[] {
  ensureSchema();
  const cap = Math.max(1, Math.floor(msgBudget));
  // Pull newest-first via SQL — bounded by msgBudget so we never load
  // more rows than we could possibly keep.
  const rows = db().prepare(`
    SELECT * FROM chat_messages WHERE session_id = ?
    ORDER BY ts DESC LIMIT ?
  `).all(session_id, cap) as MessageRow[];
  const chrono = rows.map(rowToMessage).reverse();

  // Group messages so that an assistant `tool_use` message and its
  // paired user `tool_result` reply are an indivisible unit. Without
  // this, the token-budget walk can keep the oversized tool_result
  // but drop the tool_use that produced it — leading to the
  // tool_result being treated as an orphan and stripped, leaving an
  // empty history. Each group is chronological.
  const groups: Message[][] = [];
  for (let i = 0; i < chrono.length; i++) {
    const m = chrono[i];
    const hasToolUse = m.role === 'assistant' && m.blocks.some((b) => b.type === 'tool_use');
    const next = chrono[i + 1];
    const nextHasToolResult =
      next && next.role === 'user' && next.blocks.some((b) => b.type === 'tool_result');
    if (hasToolUse && nextHasToolResult) {
      groups.push([m, next]);
      i++; // skip the partner — we just consumed it
    } else {
      groups.push([m]);
    }
  }

  // Walk groups newest-first, applying the token budget. Always keep
  // at least one group so an oversized last group doesn't strand the
  // loop (provider will see a single message — still valid).
  const keptGroups: Message[][] = [];
  let used = 0;
  let evictedCount = 0;
  for (let i = groups.length - 1; i >= 0; i--) {
    const g = groups[i];
    const cost = g.reduce((s, m) => s + estimateTokens(m), 0);
    if (keptGroups.length > 0 && used + cost > tokenBudget) {
      evictedCount = i + 1; // groups [0..i] would have been evicted
      break;
    }
    keptGroups.unshift(g);
    used += cost;
  }

  // ── Pin the SESSION's first user message (task brief) ──────────
  // Even if eviction would normally drop it, the user's opening prompt
  // defines the task. Losing it causes the model to lose track of
  // what was asked — symptom: model writes "summarize all X" and
  // hallucinates instead of processing the specific list the user
  // gave. Re-fetch the absolute first user message, prepend if not
  // already in keptGroups. Cap its tokens so a truly enormous brief
  // can't break the call — keep first ~2k tokens.
  if (evictedCount > 0) {
    const firstUserRow = db().prepare(`
      SELECT * FROM chat_messages WHERE session_id = ? AND role = 'user'
      ORDER BY ts ASC LIMIT 1
    `).get(session_id) as MessageRow | undefined;
    if (firstUserRow) {
      const firstUserMsg = rowToMessage(firstUserRow);
      const alreadyKept = keptGroups.some((g) => g.some((m) => m.id === firstUserMsg.id));
      if (!alreadyKept) {
        // Cap to ~2000 tokens of brief (≈8KB) — tasks longer than that
        // should be split anyway; preserving the head is enough to
        // anchor the model to the original ask.
        const FIRST_BRIEF_TOKEN_CAP = 2000;
        let pinned = firstUserMsg;
        if (estimateTokens(firstUserMsg) > FIRST_BRIEF_TOKEN_CAP) {
          pinned = clipMessageToTokens(firstUserMsg, FIRST_BRIEF_TOKEN_CAP);
        }
        keptGroups.unshift([pinned]);
        console.log(`[session-cap] pinned first user message (id=${firstUserMsg.id}) — ${evictedCount} groups evicted, ${used} tokens used / ${tokenBudget} budget`);
      }
    } else {
      console.log(`[session-cap] ${evictedCount} groups evicted, no first user message found to pin`);
    }
  }
  return keptGroups.flat();
}

/** Clip a message's text content to a soft token cap. Tool blocks are
 *  preserved verbatim (they're usually small structural data); only
 *  long text blocks get a head-only truncation with a marker. */
function clipMessageToTokens(m: Message, tokenCap: number): Message {
  const charCap = tokenCap * 4; // matches estimateTokens char/4 heuristic
  const blocks = m.blocks.map((b) => {
    if (b.type === 'text' && b.text.length > charCap) {
      return { ...b, text: b.text.slice(0, charCap) + '\n\n[…task brief truncated to keep in-context]' };
    }
    return b;
  });
  return { ...m, blocks };
}

export function deleteMessage(id: string): boolean {
  ensureSchema();
  const r = db().prepare(`DELETE FROM chat_messages WHERE id = ?`).run(id);
  return r.changes > 0;
}
