import OpenAI from 'openai';
import { Logger } from '../utils/logger.js';
import { CostTracker, TokenUsage } from '../utils/cost-tracker.js';

export interface OpenAIConfig {
  apiKey: string;
  defaultModel: string;
  maxTokens: number;
  temperature?: number;
}

export interface ChatMessage {
  role: 'system' | 'user' | 'assistant';
  content: string;
}

export interface CompletionOptions {
  model?: string;
  temperature?: number;
  maxTokens?: number;
  systemPrompt?: string;
  messages?: ChatMessage[];
}

export interface CompletionResult {
  content: string;
  usage: TokenUsage;
  model: string;
  finishReason: string;
  cost: number;
}

export class OpenAIIntegration {
  private client: OpenAI;
  private config: OpenAIConfig;
  private logger: Logger;
  private costTracker?: CostTracker;

  constructor(config: OpenAIConfig) {
    this.config = config;
    this.logger = new Logger('OpenAIIntegration');
    
    if (!config.apiKey) {
      throw new Error('OpenAI API key is required');
    }

    this.client = new OpenAI({
      apiKey: config.apiKey,
    });

    this.logger.info('OpenAI integration initialized', {
      defaultModel: config.defaultModel,
      maxTokens: config.maxTokens,
    });
  }

  public setCostTracker(costTracker: CostTracker): void {
    this.costTracker = costTracker;
    this.logger.info('Cost tracker attached to OpenAI integration');
  }

  public async createCompletion(
    prompt: string,
    options: CompletionOptions = {}
  ): Promise<CompletionResult> {
    const model = options.model || this.config.defaultModel;
    const temperature = options.temperature ?? this.config.temperature ?? 0.7;
    const maxTokens = options.maxTokens || this.config.maxTokens;

    // Build messages array
    const messages: ChatMessage[] = [];
    
    if (options.systemPrompt) {
      messages.push({
        role: 'system',
        content: options.systemPrompt,
      });
    }

    if (options.messages && options.messages.length > 0) {
      messages.push(...options.messages);
    } else {
      messages.push({
        role: 'user',
        content: prompt,
      });
    }

    this.logger.debug('Creating completion', {
      model,
      temperature,
      maxTokens,
      messageCount: messages.length,
    });

    try {
      const completion = await this.client.chat.completions.create({
        model,
        messages: messages.map(msg => ({
          role: msg.role,
          content: msg.content,
        })),
        temperature,
        max_tokens: maxTokens,
      });

      const choice = completion.choices[0];
      if (!choice || !choice.message?.content) {
        throw new Error('No completion generated');
      }

      const usage: TokenUsage = {
        promptTokens: completion.usage?.prompt_tokens || 0,
        completionTokens: completion.usage?.completion_tokens || 0,
        totalTokens: completion.usage?.total_tokens || 0,
      };

      // Track usage and cost if cost tracker is available
      let cost = 0;
      if (this.costTracker) {
        await this.costTracker.trackUsage(model, usage);
        cost = await this.calculateCost(model, usage);
      }

      const result: CompletionResult = {
        content: choice.message.content,
        usage,
        model,
        finishReason: choice.finish_reason || 'unknown',
        cost,
      };

      this.logger.debug('Completion created successfully', {
        model,
        promptTokens: usage.promptTokens,
        completionTokens: usage.completionTokens,
        totalTokens: usage.totalTokens,
        cost,
        finishReason: result.finishReason,
      });

      return result;
    } catch (error) {
      this.logger.error('Failed to create completion', {
        model,
        error: error instanceof Error ? error.message : String(error),
      });
      throw error;
    }
  }

  public async createChatCompletion(
    messages: ChatMessage[],
    options: CompletionOptions = {}
  ): Promise<CompletionResult> {
    return this.createCompletion('', {
      ...options,
      messages,
    });
  }

  public async streamCompletion(
    prompt: string,
    options: CompletionOptions = {},
    onChunk?: (chunk: string) => void
  ): Promise<CompletionResult> {
    const model = options.model || this.config.defaultModel;
    const temperature = options.temperature ?? this.config.temperature ?? 0.7;
    const maxTokens = options.maxTokens || this.config.maxTokens;

    // Build messages array
    const messages: ChatMessage[] = [];
    
    if (options.systemPrompt) {
      messages.push({
        role: 'system',
        content: options.systemPrompt,
      });
    }

    if (options.messages && options.messages.length > 0) {
      messages.push(...options.messages);
    } else {
      messages.push({
        role: 'user',
        content: prompt,
      });
    }

    this.logger.debug('Creating streaming completion', {
      model,
      temperature,
      maxTokens,
      messageCount: messages.length,
    });

    try {
      const stream = await this.client.chat.completions.create({
        model,
        messages: messages.map(msg => ({
          role: msg.role,
          content: msg.content,
        })),
        temperature,
        max_tokens: maxTokens,
        stream: true,
      });

      let fullContent = '';
      let finishReason = 'unknown';
      let usage: TokenUsage = {
        promptTokens: 0,
        completionTokens: 0,
        totalTokens: 0,
      };

      for await (const chunk of stream) {
        const delta = chunk.choices[0]?.delta;
        
        if (delta?.content) {
          fullContent += delta.content;
          if (onChunk) {
            onChunk(delta.content);
          }
        }

        if (chunk.choices[0]?.finish_reason) {
          finishReason = chunk.choices[0].finish_reason;
        }

        // Note: Usage info is typically only available in the last chunk
        if (chunk.usage) {
          usage = {
            promptTokens: chunk.usage.prompt_tokens,
            completionTokens: chunk.usage.completion_tokens,
            totalTokens: chunk.usage.total_tokens,
          };
        }
      }

      // If usage wasn't provided, estimate it
      if (usage.totalTokens === 0) {
        usage = this.estimateTokenUsage(fullContent, messages);
      }

      // Track usage and cost
      let cost = 0;
      if (this.costTracker) {
        await this.costTracker.trackUsage(model, usage);
        cost = await this.calculateCost(model, usage);
      }

      const result: CompletionResult = {
        content: fullContent,
        usage,
        model,
        finishReason,
        cost,
      };

      this.logger.debug('Streaming completion finished', {
        model,
        contentLength: fullContent.length,
        promptTokens: usage.promptTokens,
        completionTokens: usage.completionTokens,
        totalTokens: usage.totalTokens,
        cost,
        finishReason,
      });

      return result;
    } catch (error) {
      this.logger.error('Failed to create streaming completion', {
        model,
        error: error instanceof Error ? error.message : String(error),
      });
      throw error;
    }
  }

  public async generateSystemPrompt(
    role: string,
    context: string,
    capabilities: string[]
  ): Promise<string> {
    const prompt = `Generate a system prompt for an AI agent with the following specifications:

Role: ${role}
Context: ${context}
Capabilities: ${capabilities.join(', ')}

The system prompt should:
1. Clearly define the agent's role and responsibilities
2. Explain how to use the available capabilities
3. Provide guidelines for interacting with other agents
4. Include specific instructions for the BZZZ P2P network
5. Be concise but comprehensive

Generate only the system prompt, without additional explanation.`;

    const result = await this.createCompletion(prompt, {
      temperature: 0.3, // Lower temperature for consistent system prompts
      maxTokens: 1000,
    });

    return result.content;
  }

  private async calculateCost(model: string, usage: TokenUsage): Promise<number> {
    // This is a simple estimation - the CostTracker has more sophisticated pricing
    const pricing = {
      'gpt-5': { prompt: 0.05 / 1000, completion: 0.15 / 1000 },
      'gpt-4': { prompt: 0.03 / 1000, completion: 0.06 / 1000 },
      'gpt-4-turbo': { prompt: 0.01 / 1000, completion: 0.03 / 1000 },
      'gpt-3.5-turbo': { prompt: 0.0005 / 1000, completion: 0.0015 / 1000 },
    };

    const modelPricing = pricing[model as keyof typeof pricing] || pricing['gpt-5'];
    
    return (usage.promptTokens * modelPricing.prompt) + 
           (usage.completionTokens * modelPricing.completion);
  }

  private estimateTokenUsage(content: string, messages: ChatMessage[]): TokenUsage {
    // Rough estimation: ~4 characters per token for English text
    const estimateTokens = (text: string): number => Math.ceil(text.length / 4);

    const promptText = messages.map(m => m.content).join(' ');
    const promptTokens = estimateTokens(promptText);
    const completionTokens = estimateTokens(content);

    return {
      promptTokens,
      completionTokens,
      totalTokens: promptTokens + completionTokens,
    };
  }

  public async testConnection(): Promise<boolean> {
    try {
      const result = await this.createCompletion('Test connection. Respond with "OK".', {
        maxTokens: 10,
        temperature: 0,
      });

      this.logger.info('OpenAI connection test successful', {
        model: this.config.defaultModel,
        response: result.content.trim(),
        tokens: result.usage.totalTokens,
      });

      return result.content.toLowerCase().includes('ok');
    } catch (error) {
      this.logger.error('OpenAI connection test failed', {
        error: error instanceof Error ? error.message : String(error),
      });
      return false;
    }
  }

  public getDefaultModel(): string {
    return this.config.defaultModel;
  }

  public updateConfig(newConfig: Partial<OpenAIConfig>): void {
    this.config = { ...this.config, ...newConfig };
    this.logger.info('OpenAI config updated', newConfig);
  }
}