import { Database } from "bun:sqlite";
import { mkdirSync } from "fs";

mkdirSync("data", { recursive: true });
const db = new Database("data/agents.db");

db.run("PRAGMA journal_mode = WAL");

db.run(`
  CREATE TABLE IF NOT EXISTS agents (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task TEXT NOT NULL,
    status TEXT NOT NULL DEFAULT 'spawning',
    chat_id TEXT NOT NULL,
    thread_msg_id INTEGER,
    created_at TEXT DEFAULT (datetime('now')),
    updated_at TEXT DEFAULT (datetime('now')),
    summary TEXT,
    error TEXT
  )
`);

db.run(`
  CREATE TABLE IF NOT EXISTS agent_logs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    agent_id INTEGER NOT NULL,
    role TEXT NOT NULL,
    content TEXT NOT NULL,
    created_at TEXT DEFAULT (datetime('now')),
    FOREIGN KEY (agent_id) REFERENCES agents(id)
  )
`);

db.run(`
  CREATE TABLE IF NOT EXISTS agent_messages (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    agent_id INTEGER NOT NULL,
    role TEXT NOT NULL,
    content TEXT NOT NULL,
    tool_use TEXT,
    created_at TEXT DEFAULT (datetime('now')),
    FOREIGN KEY (agent_id) REFERENCES agents(id)
  )
`);

export interface Agent {
  id: number;
  task: string;
  status: string;
  chat_id: string;
  thread_msg_id: number | null;
  created_at: string;
  updated_at: string;
  summary: string | null;
  error: string | null;
}

export interface AgentLog {
  id: number;
  agent_id: number;
  role: string;
  content: string;
  created_at: string;
}

// --- Agent CRUD ---

export function createAgent(task: string, chatId: string): Agent {
  const stmt = db.prepare(
    `INSERT INTO agents (task, status, chat_id) VALUES (?, 'spawning', ?)`
  );
  stmt.run(task, chatId);
  const id = Number(db.query("SELECT last_insert_rowid() as id").get()!.id);
  return getAgent(id)!;
}

export function getAgent(id: number): Agent | undefined {
  return db.query(`SELECT * FROM agents WHERE id = ?`).get(id) as
    | Agent
    | undefined;
}

export function updateAgent(
  id: number,
  updates: Partial<Pick<Agent, "status" | "thread_msg_id" | "summary" | "error">>
): void {
  const fields: string[] = ["updated_at = datetime('now')"];
  const values: unknown[] = [];

  if (updates.status !== undefined) {
    fields.push("status = ?");
    values.push(updates.status);
  }
  if (updates.thread_msg_id !== undefined) {
    fields.push("thread_msg_id = ?");
    values.push(updates.thread_msg_id);
  }
  if (updates.summary !== undefined) {
    fields.push("summary = ?");
    values.push(updates.summary);
  }
  if (updates.error !== undefined) {
    fields.push("error = ?");
    values.push(updates.error);
  }

  values.push(id);
  db.prepare(`UPDATE agents SET ${fields.join(", ")} WHERE id = ?`).run(
    ...values
  );
}

export function listAgents(chatId?: string): Agent[] {
  if (chatId) {
    return db
      .query(`SELECT * FROM agents WHERE chat_id = ? ORDER BY id DESC`)
      .all(chatId) as Agent[];
  }
  return db.query(`SELECT * FROM agents ORDER BY id DESC`).all() as Agent[];
}

export function getActiveAgents(): Agent[] {
  return db
    .query(
      `SELECT * FROM agents WHERE status IN ('spawning', 'working', 'waiting') ORDER BY id`
    )
    .all() as Agent[];
}

export function findAgentByThreadMsg(messageId: number): Agent | undefined {
  return db
    .query(`SELECT * FROM agents WHERE thread_msg_id = ?`)
    .get(messageId) as Agent | undefined;
}

// --- Logs ---

export function addLog(agentId: number, role: string, content: string): void {
  db.prepare(
    `INSERT INTO agent_logs (agent_id, role, content) VALUES (?, ?, ?)`
  ).run(agentId, role, content);
}

export function getLogs(agentId: number, limit = 20): AgentLog[] {
  return db
    .query(
      `SELECT * FROM agent_logs WHERE agent_id = ? ORDER BY id DESC LIMIT ?`
    )
    .all(agentId, limit) as AgentLog[];
}

// --- Messages (conversation history) ---

export function addMessage(
  agentId: number,
  role: string,
  content: string,
  toolUse?: string
): void {
  db.prepare(
    `INSERT INTO agent_messages (agent_id, role, content, tool_use) VALUES (?, ?, ?, ?)`
  ).run(agentId, role, content, toolUse || null);
}

export function getMessages(
  agentId: number
): Array<{ role: string; content: string; tool_use: string | null }> {
  return db
    .query(
      `SELECT role, content, tool_use FROM agent_messages WHERE agent_id = ? ORDER BY id ASC`
    )
    .all(agentId) as Array<{
    role: string;
    content: string;
    tool_use: string | null;
  }>;
}

export default db;