import { jsonBytes } from './json.js';
import fs from 'node:fs';
import { Agent } from '@earendil-works/pi-agent-core';
import { localTools } from './local-tools.js';
import { LIMITS, type Warning, type BridgeBaseMessage, type ModelRequest, type ModelResponse } from './types.js';
import { NDJsonFrameParser, StreamMessageWriter, sanitizeError } from './framing.js';
import { envelope, confluenceResult } from './validation.js';
import { createModelProvider } from './model-provider.js';
import { createConfluenceTools } from './confluence-tools.js';
import { scanArtifactsDirectory, exportArtifacts, closeArtifacts } from './artifacts.js';
import { supervisorChannel } from './supervision.js';
import { WarningCollector } from './warnings.js';

export type BridgeState = 'INIT' | 'RUNNING' | 'COLLECTING' | 'COMPLETE' | 'FAILED';
export interface BridgeOptions {
  stdin?: NodeJS.ReadableStream; stdout?: NodeJS.WritableStream;
  workDir?: string; artifactsDir?: string;
  onChildReap?: () => Promise<void>; onSetDeadline?: (ms: number) => void;
}
type Pending = { type: string; tool?: string; transfer?: string; decisions?: string[]; resolve: (x: any) => void; reject: (e: Error) => void };
export class Bridge {
  private state: BridgeState = 'INIT';
  private seq = 0;
  private seen = new Set<string>();
  private pending = new Map<string, Pending>();
  private writer: StreamMessageWriter;
  private parser = new NDJsonFrameParser();
  private input: NodeJS.ReadableStream;
  private work: string;
  private artifacts: string;
  private agent?: Agent;
  private warnings = new WarningCollector();
  private completing = false;
  private inputEnded = false;
  private timer?: NodeJS.Timeout;
  private expires = Date.now() + LIMITS.MAX_DEADLINE_MS;
  private modelCalls = 0;
  private toolCalls = 0;
  private settled = false;
  private resolve!: () => void;
  private reject!: (e: Error) => void;
  constructor(private options: BridgeOptions = {}) {
    this.input = options.stdin || process.stdin;
    this.writer = new StreamMessageWriter(options.stdout || process.stdout);
    this.work = options.workDir || '/work';
    this.artifacts = options.artifactsDir || `${this.work}/artifacts`;
  }
  nextId(): string { return `a_${++this.seq}`; }
  getState(): BridgeState { return this.state; }
  addWarning(w: Warning): void { this.warnings.add(w); }
  private arm(ms: number): void {
    this.expires = Math.min(this.expires, Date.now() + ms);
    clearTimeout(this.timer);
    this.timer = setTimeout(() => this.failRun('query_timeout'), Math.max(1, this.expires - Date.now()));
  }
  start(): Promise<void> {
    const done = new Promise<void>((resolve, reject) => { this.resolve = resolve; this.reject = reject; });
    this.arm(LIMITS.MAX_DEADLINE_MS);
    this.parser.on('data', frame => {
      try { this.inbound(frame); } catch { this.failRun('invalid_input'); }
    });
    this.parser.on('error', () => this.failRun('invalid_input'));
    this.input.on('error', () => this.failRun('connectivity_failed'));
    this.input.on('close', () => { if (this.state !== 'COMPLETE' && this.state !== 'FAILED' && !this.inputEnded) this.failRun('connectivity_failed'); });
    this.parser.on('end', () => {
      this.inputEnded = true;
      if (this.state === 'COMPLETE') this.finishEOF();
      else if (!this.completing) this.failRun('connectivity_failed');
      // A peer can close stdin synchronously after observing complete, before
      // the write callback. Wait for that write to succeed before resolving.
    });
    this.input.pipe(this.parser);
    return done;
  }
  private finishEOF(): void {
    if (this.settled) return;
    this.settled = true; this.resolve();
  }
  private inbound(frame: BridgeBaseMessage): void {
    if (this.state === 'FAILED') return;
    envelope(frame);
    if (this.seen.has(frame.id) || this.seen.size >= 512) throw new Error('Duplicate or excess IDs');
    this.seen.add(frame.id);
    if (frame.type === 'start') {
      if (this.state !== 'INIT') throw new Error('Unexpected start');
      this.state = 'RUNNING';
      this.arm(frame.payload.remaining_ms);
      this.options.onSetDeadline?.(frame.payload.remaining_ms);
      this.run(frame.payload).catch(() => this.failRun('execution_failed'));
      return;
    }
    const p = this.pending.get(frame.reply_to!);
    if (!p || p.type !== frame.type) throw new Error('Unknown or wrong response');
    if (frame.type === 'artifact_ack' && (p.transfer !== frame.payload.transfer_id || !p.decisions!.includes(frame.payload.decision))) throw new Error('Wrong artifact acknowledgement');
    if (frame.type === 'tool_response' && frame.payload.error === null) confluenceResult(frame.payload.result, p.tool!);
    this.pending.delete(frame.reply_to!);
    if (frame.payload.error) {
      const error = sanitizeError(frame.payload.error.code);
      // Model errors are terminal, while recoverable Confluence failures are
      // returned to the SDK as ordinary failed tool calls.
      p.reject(new Error(error.message));
      if (frame.type === 'model_response') this.failRun(error.code);
    } else p.resolve(frame.payload);
  }
  private async send(type: string, payload: any): Promise<void> {
    if (this.state === 'FAILED') throw new Error('Failed run');
    await this.writer.write({ v: 1, type, id: this.nextId(), payload });
  }
  private request(type: string, payload: any, response: string, transfer?: string, decisions?: string[]): Promise<any> {
    if (this.state === 'FAILED') return Promise.reject(new Error('Failed run'));
    const id = this.nextId();
    return new Promise((resolve, reject) => {
      this.pending.set(id, { type: response, tool: payload.tool, transfer, decisions, resolve, reject });
      this.writer.write({ v: 1, type, id, payload }).catch(() => this.failRun('connectivity_failed'));
    });
  }
  async sendToolRequest(tool: string, parameters: Record<string, any>): Promise<any> {
    if (this.state !== 'RUNNING' || this.pending.size >= 4 || ++this.toolCalls > 100) throw new Error('Tool call limit');
    const payload = { tool, parameters };
    if (jsonBytes(payload, LIMITS.TOOL_REQUEST_PAYLOAD_MAX_BYTES) > LIMITS.TOOL_REQUEST_PAYLOAD_MAX_BYTES) throw new Error('Tool payload limit');
    return (await this.request('tool_request', payload, 'tool_response')).result;
  }
  async sendModelRequest(request: ModelRequest): Promise<ModelResponse> {
    if (this.state !== 'RUNNING' || this.pending.size >= 4 || ++this.modelCalls > 50) { this.failRun('execution_failed'); throw new Error('Model call limit'); }
    if (jsonBytes(request, LIMITS.MODEL_REQUEST_PAYLOAD_MAX_BYTES) > LIMITS.MODEL_REQUEST_PAYLOAD_MAX_BYTES) { this.failRun('model_context_exceeded'); throw new Error('Model payload limit'); }
    return (await this.request('model_request', request, 'model_response')).result;
  }
  private async run(start: any): Promise<void> {
    fs.mkdirSync(this.artifacts, { recursive: true });
    const { model, streamFn } = createModelProvider(start.model, async request => {
      const response = await this.sendModelRequest(request);
      if (response.usage.input_tokens > start.model.context_window_tokens) { this.failRun('model_context_exceeded'); throw new Error('Context limit'); }
      if (response.stop_reason === 'length' || response.usage.output_tokens > start.model.max_output_tokens) { this.failRun('model_output_limit'); throw new Error('Output limit'); }
      return response;
    });
    this.agent = new Agent({ initialState: {
      systemPrompt: start.system_instruction, model,
      tools: [...createConfluenceTools((name, args) => this.sendToolRequest(name, args)), ...localTools(this.work)],
    }, streamFn, toolExecution: 'sequential' });
    await this.agent.prompt(start.prompt);
    await this.agent.waitForIdle();
    if (this.state !== 'RUNNING') return;
    if (this.pending.size) throw new Error('Unanswered calls');
    const final = this.agent.state.messages.at(-1);
    if (!final || final.role !== 'assistant' || final.stopReason !== 'stop') throw new Error('Invalid final state');
    const markdown = final.content.filter(x => x.type === 'text').map(x => x.text).join('');
    if (!markdown.trim()) throw new Error('Empty final answer');
    if (Buffer.byteLength(markdown) > LIMITS.FINAL_MARKDOWN_MAX_BYTES) throw new Error('Answer too large');
    this.agent.clearAllQueues(); this.agent.abort();
    // No host /proc fallback: collection requires a supervisor or an explicit
    // test hook. Production runs always use the immutable PID-1 supervisor.
    if (!this.options.onChildReap) throw new Error('Missing process supervisor');
    await this.options.onChildReap();
    if (this.state !== 'RUNNING') return;
    const scan = scanArtifactsDirectory(this.artifacts);
    try {
      scan.warnings.forEach(w => this.addWarning(w));
      this.state = 'COLLECTING';
      await this.request('collection_start', { markdown, warnings: this.warnings.snapshot() }, 'collection_ready');
      const count = await exportArtifacts(scan.candidates, {
        sendBegin: async (transfer_id, name, size_bytes) => {
          const ack = await this.request('artifact_begin', { transfer_id, name, size_bytes }, 'artifact_ack', transfer_id, ['accept', 'skip']);
          return ack.decision;
        },
        sendChunk: async (transfer_id, index, data_base64) => this.send('artifact_chunk', { transfer_id, index, data_base64 }),
        sendEnd: async (transfer_id, size_bytes, chunks) => {
          await this.request('artifact_end', { transfer_id, size_bytes, chunks }, 'artifact_ack', transfer_id, ['stored']);
        },
      });
      if (this.getState() === 'FAILED') return;
      await this.writer.write({ v: 1, type: 'complete', id: this.nextId(), payload: { accepted_transfer_count: count } }, () => { this.completing = true; });
      if (this.getState() === 'FAILED') return;
      this.state = 'COMPLETE';
      clearTimeout(this.timer);
      if (this.inputEnded) this.finishEOF();
    } finally { closeArtifacts(scan.candidates); }
  }
  async flushOutput(timeoutMs = 1000): Promise<boolean> {
    let timer: NodeJS.Timeout | undefined;
    try {
      return await Promise.race([
        this.writer.close().then(() => true, () => false),
        new Promise<boolean>(resolve => { timer = setTimeout(() => resolve(false), timeoutMs); }),
      ]);
    } finally { clearTimeout(timer); }
  }
  failRun(code: string, _raw?: unknown): void {
    if (this.state === 'FAILED' || this.state === 'COMPLETE' || this.settled) return;
    this.state = 'FAILED'; clearTimeout(this.timer);
    this.agent?.clearAllQueues(); this.agent?.abort();
    const error = sanitizeError(code);
    for (const p of this.pending.values()) p.reject(new Error(error.message));
    this.pending.clear();
    // Failure settlement does not depend on a peer continuing to drain stdout.
    this.writer.writeTerminal({ v: 1, type: 'error', id: this.nextId(), payload: error }).catch(() => {});
    this.settled = true; this.reject(new Error(error.message));
  }
}
if (/\bbridge\.(ts|js)$/.test(process.argv[1] || '')) {
  const channel = supervisorChannel();
  const bridge = new Bridge({ workDir: process.env.AGENT_WORK_DIR || '/work',
    onChildReap: channel?.reap, onSetDeadline: channel?.deadline });
  bridge.start().then(() => process.exit(0)).catch(async () => {
    // Flush complete NDJSON lines within a bounded budget. The supervisor's
    // independent deadline still owns the maximum container lifetime.
    await bridge.flushOutput();
    process.exit(1);
  });
}