src/lib/agent.ts

import { generateText } from 'ai';
import { allTools, getToolByName } from './tools';
import { z } from 'zod';
import {
  Message,
  AgentContext,
  AgentConfig,
  AgentType,
  PromptLog
} from './types/agent';

export class AIAgent {
  public contexts: Map<string, AgentContext> = new Map();
  private model: any;
  public config: AgentConfig;
  private promptLogs: PromptLog[] = [];

  constructor(model: any, config: AgentConfig) {
    this.model = model;
    this.config = config;
    console.log(`AI Agent initialized: ${config.name} (${config.type}:${config.personality}) with tools:`, config.availableTools);
  }

  getConfig(): AgentConfig {
    return { ...this.config };
  }

  updateConfig(newConfig: AgentConfig): void {
    this.config = { ...newConfig };
    this.config.updatedAt = new Date().toISOString();
    console.log(`Agent ${this.config.id} config updated`);
  }

  private getSystemPrompt(): string {
    return this.config.systemPrompt;
  }

  getOrCreateContext(conversationId: string): AgentContext {
    if (!this.contexts.has(conversationId)) {
      this.contexts.set(conversationId, {
        conversationId,
        messages: [
          {
            role: 'system',
            content: this.getSystemPrompt(),
            timestamp: new Date().toISOString()
          }
        ],
        availableTools: this.config.availableTools
      });
    }
    return this.contexts.get(conversationId)!;
  }

  async chat(conversationId: string, userMessage: string): Promise<string> {
    const context = this.getOrCreateContext(conversationId);
    const timestamp = new Date().toISOString();

    // Add user message to context
    context.messages.push({
      role: 'user',
      content: userMessage,
      timestamp
    });

    let toolsUsed: string[] = [];

    try {
      // Generate response using the AI model
      const { text } = await generateText({
        model: this.model,
        messages: context.messages.map(msg => ({
          role: msg.role,
          content: msg.content
        }))
      });

      // Check if the response indicates tool usage
      let finalResponse = text;

      try {
        const parsed = JSON.parse(text);
        if (parsed.action === 'use_tool') {
          toolsUsed.push(parsed.tool);

          // Execute the tool
          const toolResult = await this.executeTool(parsed.tool, parsed.parameters);

          // Generate a follow-up response incorporating the tool result
          const followUpMessages = [
            ...context.messages,
            { role: 'assistant' as const, content: text },
            { role: 'user' as const, content: `Tool result: ${toolResult}. Please provide a helpful response to the user based on this information.` }
          ];

          const { text: followUpText } = await generateText({
            model: this.model,
            messages: followUpMessages
          });

          finalResponse = followUpText;
        }
      } catch {
        // If parsing fails, treat as regular response
      }

      // Add assistant response to context
      context.messages.push({
        role: 'assistant',
        content: finalResponse,
        timestamp: new Date().toISOString()
      });

      // Log the interaction for future persistence
      this.logPrompt({
        agentId: this.config.id,
        conversationId,
        timestamp,
        promptVersion: '1.0', // TODO: Implement versioning
        userMessage,
        agentResponse: finalResponse,
        toolsUsed,
        metadata: {
          agentType: this.config.type,
          personality: this.config.personality
        }
      });

      return finalResponse;

    } catch (error) {
      console.error(`Agent ${this.config.id} chat error:`, error);
      return 'Sorry, I encountered an error while processing your request. Please try again.';
    }
  }

  private async executeTool(toolName: string, parameters: any): Promise<string> {
    // Check if tool is available for this agent
    if (!this.config.availableTools.includes(toolName)) {
      return `Error: Tool "${toolName}" is not available for this agent. Available tools: ${this.config.availableTools.join(', ')}`;
    }

    const tool = getToolByName(toolName);

    if (!tool) {
      return `Error: Tool "${toolName}" not found. Available tools: ${this.config.availableTools.join(', ')}`;
    }

    try {
      // Validate parameters
      const validatedParams = tool.parameters.parse(parameters);

      // Execute the tool
      const result = await tool.execute(validatedParams);
      return result;
    } catch (error) {
      return `Error executing tool "${toolName}": ${error}`;
    }
  }

  getConversationHistory(conversationId: string): Message[] {
    const context = this.contexts.get(conversationId);
    return context ? context.messages.filter(msg => msg.role !== 'system') : [];
  }

  clearConversation(conversationId: string): void {
    this.contexts.delete(conversationId);
  }

  listAvailableTools(): Array<{name: string, description: string}> {
    return allTools
      .filter(tool => this.config.availableTools.includes(tool.name))
      .map(tool => ({
        name: tool.name,
        description: tool.description
      }));
  }

  // New methods for logging and persistence preparation
  private logPrompt(log: PromptLog): void {
    this.promptLogs.push(log);
    // In the future, this could send to Langfuse or other persistence layer
    console.log(`[${this.config.name}] Logged interaction: ${log.conversationId}`);
  }

  getPromptLogs(conversationId?: string): PromptLog[] {
    if (conversationId) {
      return this.promptLogs.filter(log => log.conversationId === conversationId);
    }
    return [...this.promptLogs];
  }

  getAgentInfo(): {
    id: string;
    type: AgentType;
    personality: string;
    name: string;
    description: string;
    isActive: boolean;
    conversationCount: number;
    lastUsed: string;
    availableTools: string[];
  } {
    const conversationCount = this.contexts.size;
    const lastLog = this.promptLogs[this.promptLogs.length - 1];

    return {
      id: this.config.id,
      type: this.config.type,
      personality: this.config.personality as string,
      name: this.config.name,
      description: this.config.description,
      isActive: true,
      conversationCount,
      lastUsed: lastLog?.timestamp || this.config.createdAt,
      availableTools: this.config.availableTools
    };
  }

  // Clean up inactive conversations (optional utility)
  cleanupOldConversations(maxAgeMs: number = 24 * 60 * 60 * 1000): void {
    const cutoff = Date.now() - maxAgeMs;
    const toDelete: string[] = [];

    this.contexts.forEach((context, conversationId) => {
      const lastMessage = context.messages[context.messages.length - 1];
      if (lastMessage?.timestamp) {
        const messageTime = new Date(lastMessage.timestamp).getTime();
        if (messageTime < cutoff) {
          toDelete.push(conversationId);
        }
      }
    });

    toDelete.forEach(id => {
      this.contexts.delete(id);
    });

    if (toDelete.length > 0) {
      console.log(`Agent ${this.config.id} cleaned up ${toDelete.length} old conversations`);
    }
  }
}