import type { Sql } from '../db.js';

export interface ToolTrace {
  id: string;
  session_id: string;
  chat_id: string;
  message_id: string | null;
  turn_number: number;
  tool_name: string;
  tool_input: unknown;
  tool_output: string | null;
  started_at: string;
  finished_at: string | null;
  latency_ms: number | null;
  tokens_used: number | null;
  cache_tokens: number | null;
  reasoning_tokens: number | null;
  error: string | null;
  outcome: string | null;
  created_at: string;
}

export interface ToolTraceInsert {
  session_id: string;
  chat_id: string;
  message_id: string | null;
  turn_number: number;
  tool_name: string;
  tool_input: unknown;
  outcome?: string;
}

export interface ToolTraceUpdate {
  finished_at?: string;
  latency_ms?: number;
  tool_output?: string;
  tokens_used?: number;
  cache_tokens?: number;
  reasoning_tokens?: number;
  error?: string;
  outcome?: string;
}

export async function insertToolTrace(
  sql: Sql,
  insert: ToolTraceInsert,
): Promise<ToolTrace> {
  const [row] = await sql<ToolTrace[]>`
    INSERT INTO tool_traces (
      session_id, chat_id, message_id, turn_number,
      tool_name, tool_input, outcome
    ) VALUES (
      ${insert.session_id}, ${insert.chat_id}, ${insert.message_id},
      ${insert.turn_number}, ${insert.tool_name},
      ${sql.json(insert.tool_input as never)},
      ${insert.outcome ?? null}
    )
    RETURNING *
  `;
  if (!row) throw new Error('insertToolTrace returned no row');
  return row;
}

export async function updateToolTrace(
  sql: Sql,
  id: string,
  updates: ToolTraceUpdate,
): Promise<ToolTrace | null> {
  const cols: string[] = [];
  const vals: any[] = [];

  if (updates.finished_at !== undefined) { cols.push('finished_at'); vals.push(updates.finished_at); }
  if (updates.latency_ms !== undefined) { cols.push('latency_ms'); vals.push(updates.latency_ms); }
  if (updates.tool_output !== undefined) { cols.push('tool_output'); vals.push(updates.tool_output); }
  if (updates.tokens_used !== undefined) { cols.push('tokens_used'); vals.push(updates.tokens_used); }
  if (updates.cache_tokens !== undefined) { cols.push('cache_tokens'); vals.push(updates.cache_tokens); }
  if (updates.reasoning_tokens !== undefined) { cols.push('reasoning_tokens'); vals.push(updates.reasoning_tokens); }
  if (updates.error !== undefined) { cols.push('error'); vals.push(updates.error); }
  if (updates.outcome !== undefined) { cols.push('outcome'); vals.push(updates.outcome); }

  if (cols.length === 0) {
    const [row] = await sql<ToolTrace[]>`SELECT * FROM tool_traces WHERE id = ${id}`;
    return row ?? null;
  }

  const setClause = cols.map((c, i) => `${c} = $${i + 1}`).join(', ');
  const [row] = await sql.unsafe<ToolTrace[]>(
    `UPDATE tool_traces SET ${setClause} WHERE id = $${cols.length + 1} RETURNING *`,
    [...vals, id],
  );
  return row ?? null;
}