/**
 * Simple Agent Definitions
 * Works with stored API keys from environment variables
 * Supports local models (Ollama, LM Studio) and tools
 */

import { executeLocalModel, LocalModelConfig } from './local-model-provider.js';
import { getTool, listTools, executeTool, formatToolsForOpenAI } from './tools.js';

export interface AgentExecutor {
  execute(input: any): Promise<any>
}

/**
 * Test Agent (No AI)
 */
export const testAgent: AgentExecutor = {
  async execute(input: any) {
    const sleepMs = typeof input?.sleep_ms === 'number' ? input.sleep_ms : 0;
    if (sleepMs > 0) {
      await new Promise((resolve) => setTimeout(resolve, sleepMs));
    }
    return {
      message: 'Hello from test agent!',
      input: input,
      timestamp: new Date().toISOString(),
      status: 'success'
    }
  }
}

/**
 * Echo Agent (No AI)
 */
export const echoAgent: AgentExecutor = {
  async execute(input: any) {
    return {
      echo: input,
      timestamp: new Date().toISOString()
    }
  }
}

/**
 * GPT-4 Assistant
 * Uses OpenAI API directly with stored credentials
 */
export const gpt4Assistant: AgentExecutor = {
  async execute(input: any) {
    const { 
      prompt, 
      systemPrompt, 
      temperature, 
      maxTokens, 
      conversationHistory,
      topP,
      frequencyPenalty,
      presencePenalty,
      tools: toolNames,
      useLocalModel,
      localModelProvider,
      localModelName,
      localModelUrl
    } = input.data || {}
    
    if (!prompt && (!conversationHistory || conversationHistory.length === 0)) {
      throw new Error('prompt is required')
    }
    
    // Check if using local model
    if (useLocalModel && localModelProvider && localModelName) {
      return executeLocalModelAgent({
        prompt,
        systemPrompt,
        temperature,
        maxTokens,
        conversationHistory,
        topP,
        frequencyPenalty,
        presencePenalty,
        tools: toolNames,
        localModelProvider,
        localModelName,
        localModelUrl
      });
    }
    
    // Remote model (OpenAI)
    const apiKey = process.env['OPENAI_API_KEY']
    if (!apiKey) {
      throw new Error('OpenAI API key not configured. Run "config" in the OS terminal.')
    }
    
    // Build messages array with conversation history
    const messages: any[] = []
    
    // Add system prompt if provided
    if (systemPrompt) {
      messages.push({ role: 'system', content: systemPrompt })
    }
    
    // Add conversation history if provided
    if (conversationHistory && Array.isArray(conversationHistory)) {
      messages.push(...conversationHistory)
    }
    
    // Add current prompt
    if (prompt) {
      messages.push({ role: 'user', content: prompt })
    }
    
    // Get tools if specified
    const tools: any[] = [];
    const toolMap: Record<string, any> = {};
    if (toolNames && Array.isArray(toolNames) && toolNames.length > 0) {
      for (const toolName of toolNames) {
        const tool = getTool(toolName);
        if (tool) {
          tools.push(tool);
          toolMap[tool.name] = tool;
        }
      }
    }
    
    const maxCompletionTokens = maxTokens ? Math.max(1, Math.min(maxTokens, 4096)) : 1000
    
    // Prepare request body
    const requestBody: any = {
      model: 'gpt-4',
      messages,
      temperature: temperature ?? 0.7,
      max_tokens: maxCompletionTokens
    };
    
    // Add advanced parameters if provided
    if (topP !== undefined) requestBody.top_p = topP;
    if (frequencyPenalty !== undefined) requestBody.frequency_penalty = frequencyPenalty;
    if (presencePenalty !== undefined) requestBody.presence_penalty = presencePenalty;
    
    // Add tools if available (OpenAI function calling)
    if (tools.length > 0) {
      requestBody.tools = formatToolsForOpenAI(tools);
      requestBody.tool_choice = 'auto'; // Let model decide when to use tools
    }
    
    const response = await fetch('https://api.openai.com/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${apiKey}`
      },
      body: JSON.stringify(requestBody)
    })
    
    if (!response.ok) {
      const error = await response.text()
      throw new Error(`OpenAI API error: ${response.status} - ${error}`)
    }
    
    const data: any = await response.json()
    const usage = data.usage || {}
    const choice = data.choices?.[0]
    
    // Handle tool calls if present
    if (choice?.message?.tool_calls && toolMap) {
      // Execute tool calls and get results
      const toolResults: any[] = [];
      for (const toolCall of choice.message.tool_calls) {
        const toolName = toolCall.function.name;
        const toolArgs = JSON.parse(toolCall.function.arguments || '{}');
        const tool = toolMap[toolName];
        
        if (tool) {
          try {
            const result = await executeTool(toolName, toolArgs);
            toolResults.push({
              tool_call_id: toolCall.id,
              role: 'tool',
              name: toolName,
              content: JSON.stringify(result)
            });
          } catch (error: any) {
            toolResults.push({
              tool_call_id: toolCall.id,
              role: 'tool',
              name: toolName,
              content: JSON.stringify({ error: error.message })
            });
          }
        }
      }
      
      // If we have tool results, make a follow-up call with tool results
      if (toolResults.length > 0) {
        const followUpMessages = [...messages, choice.message, ...toolResults];
        
        const followUpResponse = await fetch('https://api.openai.com/v1/chat/completions', {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${apiKey}`
          },
          body: JSON.stringify({
            model: 'gpt-4',
            messages: followUpMessages,
            temperature: temperature ?? 0.7,
            max_tokens: maxCompletionTokens
          })
        });
        
        if (followUpResponse.ok) {
          const followUpData: any = await followUpResponse.json();
          const followUpUsage = followUpData.usage || {};
          return {
            result: followUpData.choices?.[0]?.message?.content || 'No response',
            model: followUpData.model,
            usage: {
              totalTokens: (usage.total_tokens || 0) + (followUpUsage.total_tokens || 0),
              promptTokens: (usage.prompt_tokens || 0) + (followUpUsage.prompt_tokens || 0),
              completionTokens: (usage.completion_tokens || 0) + (followUpUsage.completion_tokens || 0)
            },
            cost: calculateCost('gpt-4', (usage.total_tokens || 0) + (followUpUsage.total_tokens || 0)),
            toolsUsed: toolResults.map(tr => tr.name)
          };
        }
      }
    }
    
    return {
      result: choice?.message?.content || 'No response',
      model: data.model,
      usage: {
        totalTokens: usage.total_tokens || 0,
        promptTokens: usage.prompt_tokens || 0,
        completionTokens: usage.completion_tokens || 0
      },
      cost: calculateCost('gpt-4', usage.total_tokens || 0),
      toolsUsed: tools.length > 0 ? tools.map(t => t.name) : undefined
    }
  }
}

/**
 * GPT-3.5 Turbo Assistant (Faster, Cheaper)
 */
export const gpt35Assistant: AgentExecutor = {
  async execute(input: any) {
    const { 
      prompt, 
      systemPrompt, 
      temperature, 
      maxTokens, 
      conversationHistory,
      topP,
      frequencyPenalty,
      presencePenalty,
      tools: toolNames,
      useLocalModel,
      localModelProvider,
      localModelName,
      localModelUrl
    } = input.data || {}
    
    if (!prompt && (!conversationHistory || conversationHistory.length === 0)) {
      throw new Error('prompt is required')
    }
    
    // Check if using local model
    if (useLocalModel && localModelProvider && localModelName) {
      return executeLocalModelAgent({
        prompt,
        systemPrompt,
        temperature,
        maxTokens,
        conversationHistory,
        topP,
        frequencyPenalty,
        presencePenalty,
        tools: toolNames,
        localModelProvider,
        localModelName,
        localModelUrl
      });
    }
    
    const apiKey = process.env['OPENAI_API_KEY']
    if (!apiKey) {
      throw new Error('OpenAI API key not configured. Run "config" in the OS terminal.')
    }
    
    // Build messages array with conversation history
    const messages: any[] = []
    
    if (systemPrompt) {
      messages.push({ role: 'system', content: systemPrompt })
    }
    
    if (conversationHistory && Array.isArray(conversationHistory)) {
      messages.push(...conversationHistory)
    }
    
    if (prompt) {
      messages.push({ role: 'user', content: prompt })
    }
    
    // Get tools if specified
    const tools: any[] = [];
    const toolMap: Record<string, any> = {};
    if (toolNames && Array.isArray(toolNames) && toolNames.length > 0) {
      for (const toolName of toolNames) {
        const tool = getTool(toolName);
        if (tool) {
          tools.push(tool);
          toolMap[tool.name] = tool;
        }
      }
    }
    
    const maxCompletionTokens = maxTokens ? Math.max(1, Math.min(maxTokens, 4096)) : 1000
    
    const requestBody: any = {
      model: 'gpt-3.5-turbo',
      messages,
      temperature: temperature ?? 0.7,
      max_tokens: maxCompletionTokens
    };
    
    if (topP !== undefined) requestBody.top_p = topP;
    if (frequencyPenalty !== undefined) requestBody.frequency_penalty = frequencyPenalty;
    if (presencePenalty !== undefined) requestBody.presence_penalty = presencePenalty;
    
    if (tools.length > 0) {
      requestBody.tools = formatToolsForOpenAI(tools);
      requestBody.tool_choice = 'auto';
    }
    
    const response = await fetch('https://api.openai.com/v1/chat/completions', {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${apiKey}`
      },
      body: JSON.stringify(requestBody)
    })
    
    if (!response.ok) {
      const error = await response.text()
      throw new Error(`OpenAI API error: ${response.status} - ${error}`)
    }
    
    const data: any = await response.json()
    const usage = data.usage || {}
    const choice = data.choices?.[0]
    
    // Handle tool calls (same as gpt4Assistant)
    if (choice?.message?.tool_calls && toolMap) {
      const toolResults: any[] = [];
      for (const toolCall of choice.message.tool_calls) {
        const toolName = toolCall.function.name;
        const toolArgs = JSON.parse(toolCall.function.arguments || '{}');
        const tool = toolMap[toolName];
        
        if (tool) {
          try {
            const result = await executeTool(toolName, toolArgs);
            toolResults.push({
              tool_call_id: toolCall.id,
              role: 'tool',
              name: toolName,
              content: JSON.stringify(result)
            });
          } catch (error: any) {
            toolResults.push({
              tool_call_id: toolCall.id,
              role: 'tool',
              name: toolName,
              content: JSON.stringify({ error: error.message })
            });
          }
        }
      }
      
      if (toolResults.length > 0) {
        const followUpMessages = [...messages, choice.message, ...toolResults];
        
        const followUpResponse = await fetch('https://api.openai.com/v1/chat/completions', {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${apiKey}`
          },
          body: JSON.stringify({
            model: 'gpt-3.5-turbo',
            messages: followUpMessages,
            temperature: temperature ?? 0.7,
            max_tokens: maxCompletionTokens
          })
        });
        
        if (followUpResponse.ok) {
          const followUpData: any = await followUpResponse.json();
          const followUpUsage = followUpData.usage || {};
          return {
            result: followUpData.choices?.[0]?.message?.content || 'No response',
            model: followUpData.model,
            usage: {
              totalTokens: (usage.total_tokens || 0) + (followUpUsage.total_tokens || 0),
              promptTokens: (usage.prompt_tokens || 0) + (followUpUsage.prompt_tokens || 0),
              completionTokens: (usage.completion_tokens || 0) + (followUpUsage.completion_tokens || 0)
            },
            cost: calculateCost('gpt-3.5-turbo', (usage.total_tokens || 0) + (followUpUsage.total_tokens || 0)),
            toolsUsed: toolResults.map(tr => tr.name)
          };
        }
      }
    }
    
    return {
      result: choice?.message?.content || 'No response',
      model: data.model,
      usage: {
        totalTokens: usage.total_tokens || 0,
        promptTokens: usage.prompt_tokens || 0,
        completionTokens: usage.completion_tokens || 0
      },
      cost: calculateCost('gpt-3.5-turbo', usage.total_tokens || 0),
      toolsUsed: tools.length > 0 ? tools.map(t => t.name) : undefined
    }
  }
}

/**
 * Execute local model agent
 */
async function executeLocalModelAgent(params: {
  prompt?: string;
  systemPrompt?: string;
  temperature?: number;
  maxTokens?: number;
  conversationHistory?: Array<{ role: string; content: string }>;
  topP?: number;
  frequencyPenalty?: number;
  presencePenalty?: number;
  tools?: string[];
  localModelProvider: 'ollama' | 'lm-studio' | 'custom';
  localModelName: string;
  localModelUrl?: string;
}): Promise<any> {
  const {
    prompt,
    systemPrompt,
    temperature,
    maxTokens,
    conversationHistory,
    topP,
    frequencyPenalty,
    presencePenalty,
    localModelProvider,
    localModelName,
    localModelUrl
  } = params;
  
  // Build messages array
  const messages: Array<{ role: string; content: string }> = [];
  
  if (systemPrompt) {
    messages.push({ role: 'system', content: systemPrompt });
  }
  
  if (conversationHistory && Array.isArray(conversationHistory)) {
    messages.push(...conversationHistory);
  }
  
  if (prompt) {
    messages.push({ role: 'user', content: prompt });
  }
  
  // Configure local model (only include defined properties)
  const localConfig: LocalModelConfig = {
    provider: localModelProvider,
    model: localModelName,
    ...(localModelUrl && { baseUrl: localModelUrl }),
    ...(temperature !== undefined && { temperature }),
    ...(maxTokens !== undefined && { maxTokens }),
    ...(topP !== undefined && { topP }),
    ...(frequencyPenalty !== undefined && { frequencyPenalty }),
    ...(presencePenalty !== undefined && { presencePenalty })
  };
  
  // Execute local model
  const result = await executeLocalModel(localConfig, messages);
  
  return {
    result: result.result,
    model: result.model,
    usage: result.usage,
    cost: 0, // Local models are free
    localModel: true
  };
}

/**
 * Calculate approximate cost based on model and token usage
 */
function calculateCost(model: string, totalTokens: number): number {
  const pricing: Record<string, number> = {
    'gpt-4': 0.03 / 1000, // $0.03 per 1K tokens (approximate)
    'gpt-4-turbo': 0.01 / 1000,
    'gpt-3.5-turbo': 0.0015 / 1000
  }
  
  const pricePerToken = pricing[model] || 0.01 / 1000
  return totalTokens * pricePerToken
}

/**
 * Agent Registry
 * Maps agent names to their executors
 */
export const agentRegistry: Record<string, AgentExecutor> = {
  'test': testAgent,
  'test-agent': testAgent,
  'echo': echoAgent,
  'gpt4-assistant': gpt4Assistant,
  'gpt35-assistant': gpt35Assistant,
  'gpt-4': gpt4Assistant, // Alias
  'gpt-3.5-turbo': gpt35Assistant // Alias
}

/**
 * Get an agent by name
 */
export function getAgent(name: string): AgentExecutor | undefined {
  return agentRegistry[name]
}

