import { mkdir } from 'node:fs/promises';
import { join } from 'node:path';

const send = event => { if (process.connected) process.send(event); };
let session, running = false, disposed = false, requestNumber = 0, modelRequestNumber = 0, compaction = false, contextWindow = 0;
const pendingTools = new Map();
const schema = properties => ({ type: 'object', properties, required: Object.keys(properties).filter(key => !['timeout', 'timeout_seconds'].includes(key)), additionalProperties: false });
const string = description => ({ type: 'string', description });

export function toolSpecifications(condition, id) {
  const tools = [
    ['bash', 'Run shell commands only inside your assigned isolated container checkout. Use for reading, editing, compilation and tests. Each call starts in the checkout.', { command: string('Shell command'), timeout: { type: 'number', minimum: 1, maximum: 600 } }],
  ];
  if (condition !== 'single') tools.push(
    ['send_message', 'Send an attributed message to another persistent peer. It resumes an idle recipient.', { agent_id: string('Recipient agent ID'), message: string('Message') }],
    ['wait_for_message', 'Wait briefly for an incoming message. This tool returns the message or a timeout.', { timeout_seconds: { type: 'number', minimum: 1, maximum: 120 } }],
    ['agent_status', 'List persistent agent IDs and lifecycle state.', {}],
  );
  if (condition === 'async' && id === 'agent-0') tools.push(
    ['spawn_agent', 'Create a persistent child for this task and return immediately. The child sees your task and shared operating rules, not your original conversation. At most five agents may be active.', { task: string('Complete delegated task instructions') }],
    ['resume_agent', 'Give an existing persistent child more work without resetting its history or budget.', { agent_id: string('Agent ID'), message: string('New instructions') }],
    ['terminate_agent', 'Abort a child and permanently mark it terminated.', { agent_id: string('Child ID') }],
  );
  return tools.map(([name, description, properties]) => ({ name, label: name, description, parameters: schema(properties) }));
}

export async function createSession(meta, sdk, emit = send) {
  const { config, id, directory } = meta;
  await mkdir(directory, { recursive: true });
  const authStorage = sdk.AuthStorage.inMemory();
  authStorage.setRuntimeApiKey('study-proxy', process.env.PI_STUDY_API_KEY ?? 'study-proxy');
  const registry = sdk.ModelRegistry.inMemory(authStorage);
  const modelId = config.model_id ?? 'Beam-501B-A23B';
  registry.registerProvider('study-proxy', {
    baseUrl: `${config.proxy_url.replace(/\/$/, '')}/runs/${encodeURIComponent(config.run_id)}/agents/${encodeURIComponent(id)}/v1`,
    api: 'openai-completions',
    apiKey: '$PI_STUDY_API_KEY',
    models: [{ id: modelId, name: modelId, reasoning: true, input: ['text'], contextWindow: config.context_window ?? 262144, maxTokens: config.max_output_tokens ?? 16000,
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
      compat: { supportsStore: false, supportsDeveloperRole: true, supportsReasoningEffort: true, supportsUsageInStreaming: true, maxTokensField: 'max_tokens', reasoningEffortMap: { minimal: 'low', low: 'low', medium: 'medium', high: 'high', xhigh: 'xhigh' } },
    }],
  });
  const settings = sdk.SettingsManager.inMemory({
    compaction: { enabled: true, reserveTokens: 20000, keepRecentTokens: 12000 },
    retry: { enabled: false, provider: { maxRetries: 0, timeoutMs: 600000 } },
    enableAnalytics: false, enableInstallTelemetry: false, packages: [],
  });
  const loader = new sdk.DefaultResourceLoader({ cwd: directory, agentDir: directory, settingsManager: settings, noExtensions: true, noSkills: true, noPromptTemplates: true, noThemes: true, noContextFiles: true, systemPromptOverride: () => 'You are a coding agent reconstructing software from its observable behavior. Follow the study task and operating instructions. You have only explicitly listed tools. All shell actions execute in the isolated solver container.' });
  await loader.reload();
  const specifications = toolSpecifications(config.condition, id);
  const customTools = specifications.map(spec => ({ ...spec, execute: async (_toolCallId, args, signal) => {
    const request_id = `${id}-tool-${++requestNumber}`;
    return new Promise(resolve => {
      const abort = () => { pendingTools.delete(request_id); resolve({ content: [{ type: 'text', text: 'Tool cancelled' }], isError: true }); };
      if (signal?.aborted) { abort(); return; }
      signal?.addEventListener('abort', abort, { once: true });
      pendingTools.set(request_id, result => { signal?.removeEventListener('abort', abort); resolve(result); });
      emit({ type: 'tool_request', request_id, name: spec.name, args });
    });
  } }));
  const { session: created } = await sdk.createAgentSession({
    cwd: directory, agentDir: directory, authStorage, modelRegistry: registry,
    model: registry.find('study-proxy', modelId), thinkingLevel: config.thinking_level ?? 'medium',
    noTools: 'builtin', tools: specifications.map(tool => tool.name), customTools,
    resourceLoader: loader, settingsManager: settings,
    sessionManager: sdk.SessionManager.create(directory, join(directory, 'sessions')),
  });
  const enabled = created.getActiveToolNames();
  if (enabled.some(name => !specifications.some(spec => spec.name === name)) || enabled.length !== specifications.length) throw new Error(`Unexpected active tools: ${enabled.join(',')}`);
  const originalStream = created.agent.streamFn;
  created.agent.streamFn = async (model, context, options = {}) => {
    const request_id = `${id}-model-${++modelRequestNumber}`;
    const started = performance.now();
    emit({ type: 'model_request_start', request_id, purpose: compaction ? 'compaction' : 'agent', context });
    const priorPayload = options.onPayload;
    const stream = await originalStream(model, context, { ...options, maxTokens: config.max_output_tokens ?? 16000, reasoning: config.thinking_level ?? 'medium',
      headers: { ...options.headers, 'X-Study-Request-ID': request_id, 'X-Study-Purpose': compaction ? 'compaction' : 'agent', 'X-Study-Context-Window': `${id}-${contextWindow}` },
      onPayload: async (payload, requestModel) => {
        const before = priorPayload ? await priorPayload(payload, requestModel) : payload;
        const effective = { ...(before ?? payload), temperature: config.temperature ?? 0.7, top_p: config.top_p ?? 0.9, max_tokens: config.max_output_tokens ?? 16000, reasoning_effort: config.thinking_level ?? 'medium' };
        delete effective.max_completion_tokens;
        return effective;
      },
    });
    void stream.result().then(message => emit({ type: 'model_request_end', request_id, duration_ms: performance.now() - started, usage: message.usage, stop_reason: message.stopReason, error: message.errorMessage }), error => emit({ type: 'model_request_error', request_id, error: String(error) }));
    return stream;
  };
  created.subscribe(event => {
    if (event.type === 'compaction_start') compaction = true;
    if (event.type === 'compaction_end') { compaction = false; if (event.result && !event.aborted) contextWindow++; }
    if (event.type === 'message_start' && event.message?.role === 'user') {
      const content = typeof event.message.content === 'string' ? event.message.content : event.message.content.filter(block => block.type === 'text').map(block => block.text).join('\n');
      const delivered = content.match(/^Message from (agent-\d+) \(([^;]+); event (\d+)\):/);
      if (delivered) emit({ type: 'message_delivered', sender_id: delivered[1], kind: delivered[2], causal_event_id: Number(delivered[3]), delivery: 'context_insertion' });
    }
    // Final messages preserve usage/thinking; streaming deltas would duplicate large payloads.
    if (event.type !== 'message_update' && event.type !== 'tool_execution_update') emit({ type: 'pi_event', event });
  });
  emit({ type: 'session_created', session_id: created.sessionId, session_file: created.sessionFile, tools: enabled });
  return created;
}

async function prompt(message) {
  if (disposed) return;
  if (running) { await session.steer(message); return; }
  running = true;
  let error;
  try { await session.prompt(message, { expandPromptTemplates: false }); }
  catch (caught) { error = String(caught); }
  finally {
    running = false;
    const last = [...session.messages].reverse().find(message => message.role === 'assistant');
    if (last?.stopReason === 'error') error = last.errorMessage ?? 'Provider error';
    send({ type: 'turn_done', answer: last?.content.filter(block => block.type === 'text').map(block => block.text).join('\n') ?? '', error });
  }
}

if (process.send) {
  process.on('message', async message => {
    try {
      if (message.type === 'init') {
        const sdk = await import('@earendil-works/pi-coding-agent');
        session = await createSession(message, sdk); send({ type: 'ready' });
      } else if (message.type === 'prompt') { void prompt(message.message).catch(error => send({ type: 'worker_error', error: String(error) })); }
      else if (message.type === 'tool_result') { pendingTools.get(message.request_id)?.(message.result); pendingTools.delete(message.request_id); }
      else if (message.type === 'abort') { await session?.abort(); session?.abortCompaction(); }
      else if (message.type === 'dispose') { disposed = true; await session?.abort(); session?.abortCompaction(); session?.dispose(); process.exit(0); }
    } catch (error) { send({ type: 'worker_error', error: String(error) }); }
  });
  process.on('disconnect', () => { void session?.abort(); process.exit(0); });
}