import { AgentConfig, Message, ToolCall, StreamEvent } from '../types';
import { ProviderManager } from '../providers/index';
import { Memory, MemoryInterface } from './memory';
import { ContextBuilder } from './context';
import { SkillsLoader } from './skills';
import { ToolRegistry } from './tools/registry';
import { ShellTool } from './tools/shell';
import { ReadFileTool, WriteFileTool } from './tools/file';
import { Config } from '../config/schema';
import { logger } from '../utils/logger';

/**
 * Agent response
 */
export interface AgentResponse {
  content: string;
  toolCalls?: ToolCall[];
  finishReason?: string;
}

export interface AgentLoopOptions {
  memory?: MemoryInterface;
  systemPromptOverride?: string;
  maxIterations?: number;
  registerBuiltInTools?: boolean;
}

/**
 * Main agent loop - handles LLM and tool execution
 */
export class AgentLoop {
  private config: AgentConfig;
  private providerManager: ProviderManager;
  private memory: MemoryInterface;
  private contextBuilder: ContextBuilder;
  private skillsLoader: SkillsLoader;
  private toolRegistry: ToolRegistry;
  private maxIterations: number;

  constructor(
    sessionId: string,
    config: Config,
    agentConfig?: Partial<AgentConfig>,
    options?: AgentLoopOptions
  ) {
    this.config = {
      model: config.agents?.defaults?.model || 'anthropic/claude-opus-4-5',
      temperature: config.agents?.defaults?.temperature || 0.7,
      maxTokens: config.agents?.defaults?.maxTokens || 4096,
      systemPrompt: config.agents?.defaults?.systemPrompt,
      ...agentConfig,
    };

    if (options?.systemPromptOverride) {
      this.config.systemPrompt = options.systemPromptOverride;
    }

    this.providerManager = new ProviderManager(config);
    this.memory = options?.memory ?? new Memory(sessionId);
    this.contextBuilder = new ContextBuilder(this.config);
    this.skillsLoader = new SkillsLoader();
    this.toolRegistry = new ToolRegistry();
    this.maxIterations = options?.maxIterations ?? 10;

    if (options?.registerBuiltInTools !== false) {
      this.registerBuiltInTools(config);
    }
  }

  private registerBuiltInTools(config: Config): void {
    const toolsConfig = config.tools;

    this.toolRegistry.register(
      new ShellTool(
        toolsConfig?.restrictToWorkspace,
        toolsConfig?.allowedCommands,
        toolsConfig?.deniedCommands
      )
    );
    this.toolRegistry.register(new ReadFileTool());
    this.toolRegistry.register(new WriteFileTool());
  }

  private buildContext(): Message[] {
    const skills = this.skillsLoader.getSkills();
    const tools = this.toolRegistry.getDefinitions();
    const conversationMessages = this.memory.getMessages();
    const rawContextMessages = this.contextBuilder.buildContextMessages(
      conversationMessages,
      skills,
      tools
    );

    const CHARS_PER_TOKEN = 4;
    const CONTEXT_TO_RESPONSE_RATIO = 4;
    const maxContextChars = (this.config.maxTokens || 4096) * CHARS_PER_TOKEN * CONTEXT_TO_RESPONSE_RATIO;
    return this.contextBuilder.truncateContext(rawContextMessages, maxContextChars);
  }

  /**
   * Process a user message and generate response (non-streaming)
   */
  async processMessage(userMessage: string): Promise<AgentResponse> {
    this.memory.addMessage({
      role: 'user',
      content: userMessage,
    });

    let iteration = 0;
    let continueLoop = true;
    let finalResponse: AgentResponse | null = null;

    while (continueLoop && iteration < this.maxIterations) {
      iteration++;
      logger.debug({ iteration, maxIterations: this.maxIterations }, 'Agent loop iteration');

      try {
        const contextMessages = this.buildContext();
        const tools = this.toolRegistry.getDefinitions();

        const response = await this.providerManager.complete(
          contextMessages,
          this.config.model,
          this.config.temperature,
          this.config.maxTokens,
          tools
        );

        logger.debug(
          {
            iteration,
            hasContent: !!response.content,
            toolCallsCount: response.toolCalls?.length || 0,
            finishReason: response.finishReason,
          },
          'LLM response received'
        );

        if (response.toolCalls && response.toolCalls.length > 0) {
          this.memory.addMessage({
            role: 'assistant',
            content: response.content || '',
            tool_calls: response.toolCalls,
          });

          for (const toolCall of response.toolCalls) {
            const toolName = toolCall.function.name;
            let toolArgs: Record<string, unknown>;
            try {
              toolArgs = JSON.parse(toolCall.function.arguments) as Record<string, unknown>;
            } catch {
              logger.warn(
                { tool: toolName, arguments: toolCall.function.arguments },
                'Invalid JSON in tool arguments, skipping tool call'
              );
              this.memory.addMessage({
                role: 'tool',
                content: `Error: Invalid JSON arguments for tool ${toolName}`,
                name: toolName,
                tool_call_id: toolCall.id,
              });
              continue;
            }

            logger.info({ tool: toolName, args: toolArgs }, 'Executing tool');

            const toolResult = await this.toolRegistry.execute(toolName, toolArgs);

            this.memory.addMessage({
              role: 'tool',
              content: toolResult.success ? toolResult.output : `Error: ${toolResult.error}`,
              name: toolName,
              tool_call_id: toolCall.id,
            });
          }

          continueLoop = true;
        } else {
          this.memory.addMessage({
            role: 'assistant',
            content: response.content,
          });

          finalResponse = {
            content: response.content,
            finishReason: response.finishReason,
          };

          continueLoop = false;
        }
      } catch (error) {
        logger.error({ error, iteration }, 'Error in agent loop');
        throw error;
      }
    }

    if (iteration >= this.maxIterations) {
      logger.warn({ maxIterations: this.maxIterations }, 'Max iterations reached');
    }

    if (!finalResponse) {
      finalResponse = {
        content: 'I apologize, but I was unable to complete your request.',
        finishReason: 'max_iterations',
      };
    }

    return finalResponse;
  }

  /**
   * Process a user message with streaming
   */
  async *streamMessage(userMessage: string): AsyncGenerator<StreamEvent> {
    this.memory.addMessage({
      role: 'user',
      content: userMessage,
    });

    let iteration = 0;

    while (iteration < this.maxIterations) {
      iteration++;
      logger.debug({ iteration, maxIterations: this.maxIterations }, 'Stream loop iteration');

      try {
        const contextMessages = this.buildContext();
        const tools = this.toolRegistry.getDefinitions();

        let fullText = '';
        const accumulatedToolCalls: Map<number, { id: string; name: string; arguments: string }> = new Map();

        for await (const chunk of this.providerManager.stream(
          contextMessages,
          this.config.model,
          this.config.temperature,
          this.config.maxTokens,
          tools
        )) {
          if (chunk.type === 'delta' && chunk.content) {
            fullText += chunk.content;
            yield { type: 'text', text: chunk.content };
          }

          if (chunk.type === 'tool_calls' && chunk.toolCalls) {
            for (const tc of chunk.toolCalls) {
              const existing = accumulatedToolCalls.get(tc.index);
              if (existing) {
                if (tc.function?.arguments) {
                  existing.arguments += tc.function.arguments;
                }
              } else {
                accumulatedToolCalls.set(tc.index, {
                  id: tc.id || '',
                  name: tc.function?.name || '',
                  arguments: tc.function?.arguments || '',
                });
              }
            }
          }
        }

        if (accumulatedToolCalls.size > 0) {
          const toolCallsList: ToolCall[] = Array.from(accumulatedToolCalls.values()).map((tc) => ({
            id: tc.id,
            type: 'function' as const,
            function: { name: tc.name, arguments: tc.arguments },
          }));

          this.memory.addMessage({
            role: 'assistant',
            content: fullText,
            tool_calls: toolCallsList,
          });

          for (const toolCall of toolCallsList) {
            const toolName = toolCall.function.name;
            let toolArgs: Record<string, unknown>;
            try {
              toolArgs = JSON.parse(toolCall.function.arguments) as Record<string, unknown>;
            } catch {
              logger.warn({ tool: toolName }, 'Invalid JSON in tool arguments');
              this.memory.addMessage({
                role: 'tool',
                content: `Error: Invalid JSON arguments for tool ${toolName}`,
                name: toolName,
                tool_call_id: toolCall.id,
              });
              yield { type: 'error', error: `Invalid JSON arguments for tool ${toolName}` };
              continue;
            }

            yield { type: 'tool_use', toolUse: { id: toolCall.id, name: toolName, input: toolArgs } };

            const toolResult = await this.toolRegistry.execute(toolName, toolArgs);
            const resultContent = toolResult.success ? toolResult.output : `Error: ${toolResult.error}`;

            this.memory.addMessage({
              role: 'tool',
              content: resultContent,
              name: toolName,
              tool_call_id: toolCall.id,
            });

            yield { type: 'tool_result', text: resultContent };
          }

          continue;
        }

        this.memory.addMessage({
          role: 'assistant',
          content: fullText,
        });

        yield { type: 'done', text: fullText };
        return;
      } catch (error) {
        logger.error({ error, iteration }, 'Error in stream loop');
        yield { type: 'error', error: (error as Error).message };
        return;
      }
    }

    logger.warn({ maxIterations: this.maxIterations }, 'Max iterations reached in stream');
    yield { type: 'done', text: 'I apologize, but I was unable to complete your request.' };
  }

  getHistory(): Message[] {
    return this.memory.getMessages();
  }

  clearHistory(): void {
    this.memory.clear();
  }

  getMemory(): MemoryInterface {
    return this.memory;
  }

  getToolRegistry(): ToolRegistry {
    return this.toolRegistry;
  }

  getSkillsLoader(): SkillsLoader {
    return this.skillsLoader;
  }

  getProviderManager(): ProviderManager {
    return this.providerManager;
  }
}