import type { EngineChatToolCall } from './protocol';

export interface StreamedToolCallDelta {
  index: number;
  id?: string;
  type?: 'function';
  function?: {
    name?: string;
    arguments?: string;
  };
}

interface MutableToolCall {
  id: string;
  name: string;
  arguments: string;
}

export interface StreamedToolCallProgress {
  index: number;
  id: string;
  name: string;
  argumentCharacters: number;
}

const MAX_TOOL_CALLS = 4;
const MAX_ARGUMENT_CHARACTERS = 64 * 1024;

function assert(condition: unknown, message: string): asserts condition {
  if (!condition) {
    throw new Error(`Invalid streamed tool call: ${message}`);
  }
}

export class ToolCallAccumulator {
  private readonly calls = new Map<number, MutableToolCall>();

  hasCalls(): boolean {
    return this.calls.size > 0;
  }

  append(deltas: readonly StreamedToolCallDelta[] | undefined): StreamedToolCallProgress[] {
    if (!deltas) {
      return [];
    }
    const changed = new Set<number>();
    for (const delta of deltas) {
      assert(Number.isSafeInteger(delta.index) && delta.index >= 0, 'index must be a non-negative integer');
      assert(delta.index < MAX_TOOL_CALLS, `index ${delta.index} exceeds the ${MAX_TOOL_CALLS}-call round limit`);
      assert(delta.type === undefined || delta.type === 'function', `index ${delta.index} has an unsupported type`);
      const current = this.calls.get(delta.index) ?? { id: '', name: '', arguments: '' };
      if (delta.id !== undefined) {
        assert(typeof delta.id === 'string', `index ${delta.index} id must be a string`);
        current.id += delta.id;
      }
      if (delta.function?.name !== undefined) {
        assert(typeof delta.function.name === 'string', `index ${delta.index} name must be a string`);
        current.name += delta.function.name;
      }
      if (delta.function?.arguments !== undefined) {
        assert(typeof delta.function.arguments === 'string', `index ${delta.index} arguments must be a string`);
        current.arguments += delta.function.arguments;
        assert(current.arguments.length <= MAX_ARGUMENT_CHARACTERS, `index ${delta.index} arguments exceed 64 KiB`);
      }
      this.calls.set(delta.index, current);
      changed.add(delta.index);
    }
    return [...changed]
      .sort((left, right) => left - right)
      .map((index) => {
        const call = this.calls.get(index)!;
        return {
          index,
          id: call.id,
          name: call.name,
          argumentCharacters: call.arguments.length,
        };
      });
  }

  finish(required = false): EngineChatToolCall[] {
    if (this.calls.size === 0) {
      assert(!required, 'finish_reason was tool_calls, but the stream contained no tool calls');
      return [];
    }
    try {
      const indexes = [...this.calls.keys()].sort((left, right) => left - right);
      const result: EngineChatToolCall[] = [];
      for (let position = 0; position < indexes.length; position += 1) {
        const index = indexes[position];
        assert(index === position, `missing index ${position}`);
        const call = this.calls.get(index);
        assert(call !== undefined, `missing index ${position}`);
        assert(call.id.length > 0, `index ${index} has no id`);
        assert(call.name.length > 0, `index ${index} has no function name`);
        let parsedArguments: unknown;
        try {
          parsedArguments = JSON.parse(call.arguments);
        } catch {
          throw new Error(`Invalid streamed tool call: index ${index} arguments are not valid JSON`);
        }
        assert(
          typeof parsedArguments === 'object' && parsedArguments !== null && !Array.isArray(parsedArguments),
          `index ${index} arguments must be a JSON object`,
        );
        result.push({
          id: call.id,
          type: 'function',
          function: {
            name: call.name,
            arguments: call.arguments,
          },
        });
      }
      return result;
    } catch (error) {
      if (!required) return [];
      throw error;
    }
  }
}