import { ChatMessage, ChatOptions, ProviderResponse } from '../../../types/chat';
import { Tool } from '../../../types/tool';

export interface VertexProviderConfig {
  projectId: string;
  location: string;
  model: string;
  credentials?: any; // Google Cloud credentials
}

export interface VertexChatRequestBody {
  instances: {
    context?: string;
    examples?: Array<{
      input: { content: string };
      output: { content: string };
    }>;
    messages: Array<{
      author: string;
      content: string;
    }>;
  }[];
  parameters: {
    temperature: number;
    maxOutputTokens: number;
    topP: number;
    topK: number;
  };
}

export interface VertexChatResponse {
  predictions: Array<{
    candidates: Array<{
      author: string;
      content: string;
    }>;
    safetyAttributes?: {
      blocked: boolean;
      categories: string[];
      scores: number[];
    };
  }>;
}

export interface VertexToolSchema {
  name: string;
  description: string;
  parameters: {
    type: string;
    properties: {
      [key: string]: {
        type: string;
        description: string;
        enum?: string[];
      };
    };
    required: string[];
  };
}

export type VertexCacheKey = {
  messages: ChatMessage[];
  options: ChatOptions;
  tools?: Tool[];
};

export enum VertexModel {
  GEMINI_PRO = 'gemini-pro',
  GEMINI_PRO_VISION = 'gemini-pro-vision',
  CLAUDE_3_SONNET = 'claude-3-sonnet-20240229',
  CLAUDE_3_OPUS = 'claude-3-opus-20240229',
  // Add more models as needed
}