interface BaseMessage {
  role: 'user' | 'assistant' | 'system';
  content: ContentBlock[];
}

// Helper to create a simple text message
export const userTextMessage = (message: string): BaseMessage => {
  return {
    role: 'user',
    content: [
      {
        type: 'text',
        text: message,
      },
    ],
  };
};

type ContentBlock =
  | TextBlock
  | ImageBlock
  | ToolUseBlock
  | ToolResultBlock
  | DocumentBlock
  | ThinkingBlock
  | RedactedThinkingBlock;

interface TextBlock {
  type: 'text';
  text: string;
  [key: string]: unknown;
}

interface ImageBlock {
  type: 'image';
  source: ImageSource;
  [key: string]: unknown;
}

interface ToolUseBlock {
  type: 'toolUse';
  id: string;
  name: string;
  input: unknown;
  [key: string]: unknown;
}

interface ToolResultBlock {
  type: 'toolResult';
  toolUseId: string;
  toolName: string;
  content?: string | Array<ToolResponseBlock>;
  isError: boolean;
  [key: string]: unknown;
}

interface LLMProvider {
  name: string;
  version?: string;

  /**
   * Simple text generation with full response
   */
  generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse>;

  /**
   * Simple streaming with partial returns
   */
  stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;

  /**
   * Low-level generation API with full control
   */
  iterate(
    input: string | BaseMessage[],
    options: IterateGenerateOptions
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;
}

interface GenerateOptions {
  tools?: ToolDefinition[];
  temperature?: number;
  maxTokens?: number;
  stopSequences?: string[];
  systemPrompt?: string;
}

interface IterateGenerateOptions extends GenerateOptions {
  stream: boolean;
}

interface GenerateResponse {
  content: Array<ContentBlock>;
  tokenUsage: TokenUsage;
  duration: number;
  stopReason?: StopReason;
}

interface TokenUsage {
  inputTokens: number;
  outputTokens: number;
  inputCost: number;
  outputCost: number;
  model: string;
}

type StopReason = 'toolUse' | 'maxTokens' | 'stopSequence' | 'endTurn';

abstract class BaseLLMProvider implements LLMProvider {
  abstract name: string;
  abstract version?: string;

  abstract generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse>;

  abstract stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown>;

  async *iterate(
    input: string | BaseMessage[],
    options: IterateGenerateOptions
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown> {
    return yield* iterateFromMethods(this, input, options);
  }
}

function convertStringToMessages(
  message: string | BaseMessage[]
): BaseMessage[] {
  if (typeof message === 'string') {
    return [
      {
        role: 'user',
        content: [{ type: 'text', text: message }],
      },
    ];
  }
  return message;
}

interface ToolDefinition {
  name: string;
  description: string;
  inputSchema: JSONSchema;
}

interface JSONSchema {
  type: string;
  properties?: Record<string, unknown>;
  required?: string[];
  [key: string]: unknown;
}

import { BedrockProvider } from '@agenite/bedrock';

const provider = new BedrockProvider({
  model: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
});

// Simple text input
const response = await provider.generate('Hello, how are you?');
console.log(response.content);

// Structured message input
const response2 = await provider.generate([
  {
    role: 'user',
    content: [{ type: 'text', text: 'Tell me about AI' }],
  },
]);

const generator = provider.stream('Explain quantum computing');

for await (const chunk of generator) {
  if (chunk.type === 'text') {
    process.stdout.write(chunk.text);
  }
}

const toolDefinition = {
  name: 'calculator',
  description: 'Perform basic math operations',
  inputSchema: {
    type: 'object',
    properties: {
      operation: {
        type: 'string',
        enum: ['add', 'subtract', 'multiply', 'divide'],
      },
      a: { type: 'number' },
      b: { type: 'number' },
    },
    required: ['operation', 'a', 'b'],
  },
};

const response = await provider.generate(
  'What is 25 × 16?',
  {
    tools: [toolDefinition],
    temperature: 0,
  }
);

// Check if the response includes a tool use
const toolUse = response.content.find(
  (block) => block.type === 'toolUse'
);

if (toolUse) {
  console.log('Tool called:', toolUse.name);
  console.log('Input:', toolUse.input);
}

import { BaseLLMProvider, GenerateOptions, GenerateResponse } from '@agenite/llm';

class CustomProvider extends BaseLLMProvider {
  name = 'custom-provider';
  version = '1.0.0';
  
  constructor(private apiKey: string) {
    super();
  }
  
  async generate(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): Promise<GenerateResponse> {
    // Convert input to standard messages format
    const messages = convertStringToMessages(input);
    
    // Make API call to your custom LLM service
    // ...
    
    // Return standardized response
    return {
      content: [{ type: 'text', text: 'Response from custom LLM' }],
      tokenUsage: { model: 'custom-model', inputTokens: 10, outputTokens: 20, inputCost: 0, outputCost: 0 },
      duration: 500,
    };
  }
  
  async *stream(
    input: string | BaseMessage[],
    options?: Partial<GenerateOptions>
  ): AsyncGenerator<PartialReturn, GenerateResponse, unknown> {
    // Stream implementation
    // ...
    
    yield { type: 'text', text: 'Partial response...' };
    
    return {
      content: [{ type: 'text', text: 'Complete response' }],
      tokenUsage: { model: 'custom-model', inputTokens: 10, outputTokens: 20, inputCost: 0, outputCost: 0 },
      duration: 500,
    };
  }
}