#!/usr/bin/env node
import { randomUUID } from 'node:crypto';
import { hostname, homedir } from 'node:os';
import { readFileSync, writeFileSync, mkdirSync, unlinkSync } from 'node:fs';
import { join } from 'node:path';
import {
  type Dict,
  isRecord,
  pickString,
  pickNumber,
  truncate,
  buildEnvelope,
  createTransport,
  readStdin,
} from '../shared/lib';

const INGEST_URL = process.env.AGENTMON_INGEST_URL || 'http://localhost:8080';
const FRAMEWORK = process.env.AGENTMON_FRAMEWORK || 'claude-code';
const HOST = process.env.AGENTMON_HOST || hostname();

const { enqueue, flush } = createTransport(INGEST_URL);

// ── Persisted state (survives between hook subprocess invocations) ──────────
interface SessionState {
  runId?: string;
  spans: { [key: string]: string };   // key = sessionKey:toolName, value = spanId
  spanStartTimes?: { [spanId: string]: number }; // spanId -> epoch ms
  subagent?: { name: string; spanId: string };
  compactSpanId?: string;
}

const STATE_DIR = join(homedir(), '.agentmon-state');

function ensureStateDir() {
  try { mkdirSync(STATE_DIR, { recursive: true }); } catch { /* ignore */ }
}

function loadState(sessionKey: string): SessionState {
  try {
    const raw = readFileSync(join(STATE_DIR, sessionKey + '.json'), 'utf8');
    return JSON.parse(raw) as SessionState;
  } catch {
    return { spans: {} };
  }
}

function saveState(sessionKey: string, state: SessionState) {
  ensureStateDir();
  try {
    writeFileSync(join(STATE_DIR, sessionKey + '.json'), JSON.stringify(state), 'utf8');
  } catch { /* ignore */ }
}

function clearState(sessionKey: string) {
  try { unlinkSync(join(STATE_DIR, sessionKey + '.json')); } catch { /* ignore */ }
}
// ─────────────────────────────────────────────────────────────────────────────

const activeRuns = new Map<string, string>();
const activeSpans = new Map<string, string>();
const activeSubagents = new Map<string, { name: string; spanId: string }>();

function getSessionKey(input: Dict): string | undefined {
  return pickString(
    input.sessionId,
    input.session_id,
    input.conversationId,
    input.conversation_id,
  );
}

function getUsage(input: Dict): Dict | undefined {
  const usage = isRecord(input.usage) ? input.usage :
                isRecord(input.llm) ? input.llm :
                undefined;
  if (!usage) return undefined;

  const result: Dict = {};
  if (usage.input_tokens !== undefined) result.input_tokens = usage.input_tokens;
  if (usage.output_tokens !== undefined) result.output_tokens = usage.output_tokens;
  if (usage.cache_creation_tokens !== undefined) result.cache_creation_tokens = usage.cache_creation_tokens;
  if (usage.cache_read_tokens !== undefined) result.cache_read_tokens = usage.cache_read_tokens;
  if (usage.thinking_tokens !== undefined) result.thinking_tokens = usage.thinking_tokens;
  if (usage.total_tokens !== undefined) result.total_tokens = usage.total_tokens;
  if (usage.total_cost !== undefined) result.total_cost = usage.total_cost;

  return Object.keys(result).length > 0 ? result : undefined;
}

function getContextWindow(input: Dict): Dict | undefined {
  const stats = isRecord(input.context_stats) ? input.context_stats :
                 isRecord(input.stats) ? input.stats : undefined;
  if (!stats) return undefined;

  const result: Dict = {};
  if (stats.input_tokens !== undefined) result.input_tokens = stats.input_tokens;
  if (stats.output_tokens !== undefined) result.output_tokens = stats.output_tokens;
  if (stats.used_tokens !== undefined) result.used_tokens = stats.used_tokens;
  if (stats.max_tokens !== undefined) result.max_tokens = stats.max_tokens;
  if (stats.tokens_remaining !== undefined) result.tokens_remaining = stats.tokens_remaining;

  return Object.keys(result).length > 0 ? result : undefined;
}

function emitError(sessionKey: string | undefined, runId: string | undefined, spanId: string | undefined, errorValue: unknown) {
  if (errorValue === undefined || errorValue === null || errorValue === false) {
    return;
  }

  const errorRecord = isRecord(errorValue) ? errorValue : {};
  const message = pickString(errorRecord.message, errorRecord.error, errorValue) || 'unknown';
  const errType = pickString(errorRecord.type, errorRecord.code) || FRAMEWORK;

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'error', sessionKey, {
    runId,
    spanId,
    payload: {
      error: {
        type: errType,
        message,
      },
    },
  }));
}

async function handleSessionStart(input: Dict) {
  const sessionKey = getSessionKey(input);
  if (!sessionKey) {
    console.error('[agentmon] ignoring claude-code session.start without session_id');
    return;
  }

  const hookEventName = pickString(input.hook_event_name);
  if (hookEventName && hookEventName !== 'SessionStart') {
    console.error(`[agentmon] ignoring claude-code session.start with hook_event_name=${hookEventName}`);
    return;
  }

  const runId = randomUUID();
  activeRuns.set(sessionKey, runId);
  saveState(sessionKey, { runId, spans: {} });

  const contextWindow = getContextWindow(input);

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.start', sessionKey, {
    attributes: contextWindow ? { context_window: contextWindow } : undefined,
  }));

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.start', sessionKey, {
    runId,
    attributes: {
      trigger: pickString(input.trigger_type, input.trigger),
    },
    payload: {
      prompt_preview: truncate(input.prompt, 200),
    },
  }));

  await flush();
}

async function handleSessionEnd(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const usage = getUsage(input);
  const contextWindow = getContextWindow(input);
  const duration = pickNumber(input.duration_ms, input.elapsed_ms);

  if (runId) {
    enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.end', sessionKey, {
      runId,
      payload: {
        status: 'success',
        duration_ms: duration,
        model: pickString(input.model),
        ...(usage && { usage }),
        ...(contextWindow && { context_window: contextWindow }),
      },
    }));
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'session.end', sessionKey, {
    payload: {
      ...(usage && { usage }),
      ...(contextWindow && { context_window: contextWindow }),
    },
  }));
  activeRuns.delete(sessionKey || '');
  activeSpans.clear();
  activeSubagents.clear();
  if (sessionKey) clearState(sessionKey);

  await flush();
}

async function handlePromptSubmit(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const prompt = pickString(input.prompt, input.text, input.message);

  if (runId && prompt) {
    enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.end', sessionKey, {
      runId,
      payload: {
        status: 'success',
        duration_ms: pickNumber(input.elapsed_ms, input.duration_ms),
      },
    }));
  }

  const newRunId = randomUUID();
  if (sessionKey) {
    activeRuns.set(sessionKey, newRunId);
    saveState(sessionKey, { runId: newRunId, spans: {} });
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.start', sessionKey, {
    runId: newRunId,
    attributes: {
      type: 'user_prompt',
    },
    payload: {
      prompt_preview: truncate(prompt, 200),
    },
  }));

  await flush();
}

async function handleToolStart(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
  const toolInput = isRecord(input.input) ? input.input : {};
  const spanId = randomUUID();
  const spanKey = sessionKey ? sessionKey + ':' + toolName : undefined;

  if (spanKey) {
    activeSpans.set(spanKey, spanId);
    state.spans[spanKey] = spanId;
    state.spanStartTimes = state.spanStartTimes || {};
    state.spanStartTimes[spanId] = Date.now();
    if (runId) state.runId = runId;
    if (sessionKey) saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.start', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'tool',
      name: toolName,
    },
    payload: {
      input: truncate(toolInput, 200),
    },
  }));

  await flush();
}

async function handleToolEnd(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const toolName = pickString(input.tool, input.tool_name, input.name) || 'unknown';
  const spanKey = sessionKey ? sessionKey + ':' + toolName : undefined;
  const spanId = spanKey ? (activeSpans.get(spanKey) || state.spans[spanKey]) : undefined;
  const result = isRecord(input.result) ? input.result : isRecord(input.output) ? input.output : {};
  const success = !result.error;
  const startTime = spanId ? state.spanStartTimes?.[spanId] : undefined;
  const duration = pickNumber(input.duration_ms, input.elapsed_ms) ?? (startTime ? Date.now() - startTime : undefined);
  if (spanId && state.spanStartTimes) {
    delete state.spanStartTimes[spanId];
    if (sessionKey) saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.end', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'tool',
      name: toolName,
    },
    payload: {
      status: success ? 'success' : 'error',
      result_preview: truncate(result.output ?? result.error ?? result, 500),
      duration_ms: duration,
    },
  }));

  if (!success) {
    emitError(sessionKey, runId, spanId, result.error);
  }

  activeSpans.delete(spanKey || '');

  await flush();
}

async function handleSubagentStart(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const agentName = pickString(input.agent, input.agent_name, input.name) || 'unknown';
  const spanId = randomUUID();

  if (sessionKey) {
    activeSubagents.set(sessionKey, { name: agentName, spanId });
    state.subagent = { name: agentName, spanId };
    state.spanStartTimes = state.spanStartTimes || {};
    state.spanStartTimes[spanId] = Date.now();
    if (runId) state.runId = runId;
    saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.start', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'agent',
      name: agentName,
      type: 'subagent',
    },
    payload: {
      prompt_preview: truncate(input.prompt, 200),
    },
  }));

  await flush();
}

async function handleSubagentStop(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const subagent = sessionKey ? (activeSubagents.get(sessionKey) || state.subagent) : undefined;
  const spanId = subagent?.spanId;
  const agentName = subagent?.name || pickString(input.agent, input.agent_name) || 'unknown';
  const startTime = spanId ? state.spanStartTimes?.[spanId] : undefined;
  const duration = pickNumber(input.duration_ms, input.elapsed_ms) ?? (startTime ? Date.now() - startTime : undefined);
  const usage = getUsage(input);
  if (spanId && state.spanStartTimes) {
    delete state.spanStartTimes[spanId];
    if (sessionKey) saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.end', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'agent',
      name: agentName,
      type: 'subagent',
    },
    payload: {
      status: 'success',
      duration_ms: duration,
      ...(usage && { usage }),
    },
  }));

  activeSubagents.delete(sessionKey || '');

  await flush();
}

async function handleCompactStart(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const spanId = randomUUID();

  if (sessionKey) {
    activeSpans.set(sessionKey + ':compact', spanId);
    state.compactSpanId = spanId;
    state.spanStartTimes = state.spanStartTimes || {};
    state.spanStartTimes[spanId] = Date.now();
    if (runId) state.runId = runId;
    saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.start', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'internal',
      name: 'context_compaction',
    },
  }));

  await flush();
}

async function handleCompactEnd(input: Dict) {
  const sessionKey = getSessionKey(input);
  const state = sessionKey ? loadState(sessionKey) : { spans: {} };
  const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;
  const spanKey = sessionKey ? sessionKey + ':compact' : undefined;
  const spanId = spanKey ? (activeSpans.get(spanKey) || state.compactSpanId) : undefined;
  const startTime = spanId ? state.spanStartTimes?.[spanId] : undefined;
  const duration = pickNumber(input.duration_ms, input.elapsed_ms) ?? (startTime ? Date.now() - startTime : undefined);
  const contextWindow = getContextWindow(input);
  if (spanId && state.spanStartTimes) {
    delete state.spanStartTimes[spanId];
    if (sessionKey) saveState(sessionKey, state);
  }

  enqueue(buildEnvelope(FRAMEWORK, HOST, 'span.end', sessionKey, {
    runId,
    spanId,
    attributes: {
      span_kind: 'internal',
      name: 'context_compaction',
    },
    payload: {
      status: 'success',
      duration_ms: duration,
      ...(contextWindow && { context_window: contextWindow }),
    },
  }));

  activeSpans.delete(spanKey || '');

  await flush();
}

async function handleNotification(input: Dict) {
  const sessionKey = getSessionKey(input);
  const notificationType = pickString(input.notification_type, input.type);
  const usage = getUsage(input);
  const contextWindow = getContextWindow(input);
  const duration = pickNumber(input.duration_ms, input.elapsed_ms);

  if (notificationType === 'Done' || notificationType === 'success') {
    const state = sessionKey ? loadState(sessionKey) : { spans: {} };
    const runId = sessionKey ? (activeRuns.get(sessionKey) || state.runId) : undefined;

    if (runId) {
      enqueue(buildEnvelope(FRAMEWORK, HOST, 'run.end', sessionKey, {
        runId,
        payload: {
          status: 'success',
          duration_ms: duration,
          model: pickString(input.model),
          ...(usage && { usage }),
          ...(contextWindow && { context_window: contextWindow }),
        },
      }));
    }

    // Do NOT emit session.end or clear state here — "Done" means Claude
    // finished a turn, not the session. State must persist so tool calls
    // between turns still have a runId. The actual session.end is emitted
    // by handleSessionEnd when the Stop hook fires.
  }

  await flush();
}

const handler = async () => {
  const args = process.argv.slice(2);
  const hookType = args[0] || 'unknown';

  let input: Dict = {};
  try {
    const stdin = await readStdin();
    if (stdin) {
      input = JSON.parse(stdin);
    }
  } catch {
    input = {};
  }

  try {
    switch (hookType) {
      case 'start':
        await handleSessionStart(input);
        break;
      case 'stop':
        await handleSessionEnd(input);
        break;
      case 'prompt':
        await handlePromptSubmit(input);
        break;
      case 'tool-start':
        await handleToolStart(input);
        break;
      case 'tool-end':
        await handleToolEnd(input);
        break;
      case 'subagent-start':
        await handleSubagentStart(input);
        break;
      case 'subagent-stop':
        await handleSubagentStop(input);
        break;
      case 'compact-start':
        await handleCompactStart(input);
        break;
      case 'compact-end':
        await handleCompactEnd(input);
        break;
      case 'notification':
        await handleNotification(input);
        break;
      default:
        console.debug(`[agentmon] unknown hook type: ${hookType}`);
    }
  } catch (err) {
    console.debug('[agentmon] handler error:', err);
  }
};

handler();