/**
 * Base Agent Module
 * 
 * Provides the foundation for creating specialized agents
 * with context management, memory, and tool integration.
 */

import { randomUUID } from 'crypto';
import ZAI from 'z-ai-web-dev-sdk';
import { ConversationContextManager, CompactionConfig } from '../core/context-manager';
import { TokenCounter } from '../core/token-counter';

export interface AgentMemory {
  shortTerm: Map<string, unknown>;
  longTerm: Map<string, unknown>;
  conversationHistory: Array<{
    role: 'user' | 'assistant' | 'system';
    content: string;
    timestamp: Date;
  }>;
}

export interface AgentTool {
  name: string;
  description: string;
  execute: (params: unknown) => Promise<unknown>;
}

export interface AgentConfig {
  id?: string;
  name: string;
  description: string;
  systemPrompt: string;
  tools?: AgentTool[];
  maxTokens?: number;
  contextConfig?: Partial<CompactionConfig>;
}

export interface AgentResponse {
  content: string;
  tokens: {
    prompt: number;
    completion: number;
    total: number;
  };
  toolCalls?: Array<{
    name: string;
    params: unknown;
    result: unknown;
  }>;
  metadata?: Record<string, unknown>;
}

/**
 * BaseAgent - Foundation class for all agents
 */
export abstract class BaseAgent {
  readonly id: string;
  readonly name: string;
  readonly description: string;
  
  protected systemPrompt: string;
  protected tools: Map<string, AgentTool>;
  protected memory: AgentMemory;
  protected contextManager: ConversationContextManager;
  protected tokenCounter: TokenCounter;
  protected zai: Awaited<ReturnType<typeof ZAI.create>> | null = null;
  protected initialized = false;

  constructor(config: AgentConfig) {
    this.id = config.id || randomUUID();
    this.name = config.name;
    this.description = config.description;
    this.systemPrompt = config.systemPrompt;
    
    this.tools = new Map();
    if (config.tools) {
      for (const tool of config.tools) {
        this.tools.set(tool.name, tool);
      }
    }

    this.memory = {
      shortTerm: new Map(),
      longTerm: new Map(),
      conversationHistory: []
    };

    this.tokenCounter = new TokenCounter(config.maxTokens);
    this.contextManager = new ConversationContextManager(config.contextConfig);
  }

  /**
   * Initialize the agent
   */
  async initialize(): Promise<void> {
    if (this.initialized) return;
    this.zai = await ZAI.create();
    this.initialized = true;
  }

  /**
   * Process a user message
   */
  async process(input: string, context?: string): Promise<AgentResponse> {
    await this.initialize();

    // Add user message to context
    const userMessage = {
      role: 'user' as const,
      content: context ? `Context: ${context}\n\n${input}` : input,
      timestamp: new Date()
    };
    
    this.memory.conversationHistory.push(userMessage);
    
    // Check if context compaction is needed
    await this.contextManager.getMessages();

    // Build messages for LLM
    const messages = this.buildMessages();

    // Get response from LLM
    const response = await this.zai!.chat.completions.create({
      messages,
      thinking: { type: 'disabled' }
    });

    const assistantContent = response.choices?.[0]?.message?.content || '';

    // Add assistant response to history
    this.memory.conversationHistory.push({
      role: 'assistant',
      content: assistantContent,
      timestamp: new Date()
    });

    // Process any tool calls (if agent supports them)
    const toolCalls = await this.processToolCalls(assistantContent);

    return {
      content: assistantContent,
      tokens: {
        prompt: 0, // Would need actual token counting
        completion: 0,
        total: 0
      },
      toolCalls,
      metadata: {
        conversationLength: this.memory.conversationHistory.length
      }
    };
  }

  /**
   * Build messages array for LLM
   */
  protected buildMessages(): Array<{ role: 'user' | 'assistant' | 'system'; content: string }> {
    const messages: Array<{ role: 'user' | 'assistant' | 'system'; content: string }> = [];

    // System prompt with tool descriptions
    let systemContent = this.systemPrompt;
    if (this.tools.size > 0) {
      const toolDescriptions = Array.from(this.tools.values())
        .map(t => `- ${t.name}: ${t.description}`)
        .join('\n');
      systemContent += `\n\nAvailable tools:\n${toolDescriptions}`;
      systemContent += `\n\nTo use a tool, include [TOOL:name]params[/TOOL] in your response.`;
    }
    
    messages.push({ role: 'assistant', content: systemContent });

    // Add conversation history
    for (const msg of this.memory.conversationHistory) {
      messages.push({
        role: msg.role,
        content: msg.content
      });
    }

    return messages;
  }

  /**
   * Process tool calls in the response
   */
  protected async processToolCalls(content: string): Promise<Array<{
    name: string;
    params: unknown;
    result: unknown;
  }>> {
    const toolCalls: Array<{ name: string; params: unknown; result: unknown }> = [];
    
    // Extract tool calls from content
    const toolRegex = /\[TOOL:(\w+)\]([\s\S]*?)\[\/TOOL\]/g;
    let match;

    while ((match = toolRegex.exec(content)) !== null) {
      const toolName = match[1];
      const paramsStr = match[2].trim();
      
      const tool = this.tools.get(toolName);
      if (tool) {
        try {
          let params = paramsStr;
          try {
            params = JSON.parse(paramsStr);
          } catch {
            // Keep as string if not valid JSON
          }
          
          const result = await tool.execute(params);
          toolCalls.push({ name: toolName, params, result });
        } catch (error) {
          toolCalls.push({
            name: toolName,
            params: paramsStr,
            result: { error: String(error) }
          });
        }
      }
    }

    return toolCalls;
  }

  /**
   * Add a tool to the agent
   */
  addTool(tool: AgentTool): void {
    this.tools.set(tool.name, tool);
  }

  /**
   * Remove a tool from the agent
   */
  removeTool(name: string): boolean {
    return this.tools.delete(name);
  }

  /**
   * Store a value in short-term memory
   */
  remember(key: string, value: unknown): void {
    this.memory.shortTerm.set(key, value);
  }

  /**
   * Retrieve a value from memory
   */
  recall(key: string): unknown | undefined {
    return this.memory.shortTerm.get(key) || this.memory.longTerm.get(key);
  }

  /**
   * Store a value in long-term memory
   */
  memorize(key: string, value: unknown): void {
    this.memory.longTerm.set(key, value);
  }

  /**
   * Clear short-term memory
   */
  forget(): void {
    this.memory.shortTerm.clear();
  }

  /**
   * Clear conversation history
   */
  clearHistory(): void {
    this.memory.conversationHistory = [];
    this.contextManager.clear();
  }

  /**
   * Get conversation summary
   */
  getSummary(): string {
    const messages = this.memory.conversationHistory;
    return messages.map(m => `[${m.role}]: ${m.content.substring(0, 100)}...`).join('\n');
  }

  /**
   * Get agent stats
   */
  getStats() {
    return {
      id: this.id,
      name: this.name,
      messageCount: this.memory.conversationHistory.length,
      toolCount: this.tools.size,
      memoryItems: this.memory.shortTerm.size + this.memory.longTerm.size,
      contextStats: this.contextManager.getStats()
    };
  }

  /**
   * Abstract method for agent-specific behavior
   */
  abstract act(input: string, context?: string): Promise<AgentResponse>;
}

/**
 * SimpleAgent - A basic agent implementation
 */
export class SimpleAgent extends BaseAgent {
  async act(input: string, context?: string): Promise<AgentResponse> {
    return this.process(input, context);
  }
}

/**
 * Create a simple agent with custom system prompt
 */
export function createAgent(
  name: string,
  systemPrompt: string,
  options?: {
    description?: string;
    tools?: AgentTool[];
    maxTokens?: number;
  }
): SimpleAgent {
  return new SimpleAgent({
    name,
    systemPrompt,
    description: options?.description || `Agent: ${name}`,
    tools: options?.tools,
    maxTokens: options?.maxTokens
  });
}