import { generateText, streamText, type CoreMessage, type Tool, type StepResult } from 'ai';
import { getModel } from './client.js';
import { ToolRegistry, createDefaultRegistry } from './tools/registry.js';
import { buildSystemPrompt, type AgentMode } from './prompts/system.js';
import type { Config } from '../config.js';
import type { Logger } from '../utils/logger.js';
import type { ToolContext } from './tools/types.js';

export interface AgentRunnerOptions {
  config: Config;
  logger: Logger;
  repoRoot: string;
  model?: string;
  mode?: AgentMode;
  maxIterations?: number;
  tools?: ToolRegistry;
  onToolCall?: (name: string, input: unknown) => void;
  onToolResult?: (name: string, result: unknown) => void;
  onText?: (text: string) => void;
}

export interface ToolCallRecord {
  name: string;
  input: unknown;
  result: unknown;
}

export interface AgentResult {
  response: string;
  toolCalls: ToolCallRecord[];
  usage: { promptTokens: number; completionTokens: number };
  iterations: number;
}

/**
 * Main agent runner with AI SDK integration
 *
 * Handles the tool loop, message history, and streaming.
 */
export class AgentRunner {
  private config: Config;
  private logger: Logger;
  private repoRoot: string;
  private modelName: string;
  private mode: AgentMode;
  private maxIterations: number;
  private tools: ToolRegistry;
  private messages: CoreMessage[] = [];

  // Callbacks
  private onToolCall?: (name: string, input: unknown) => void;
  private onToolResult?: (name: string, result: unknown) => void;
  private onText?: (text: string) => void;

  constructor(options: AgentRunnerOptions) {
    this.config = options.config;
    this.logger = options.logger;
    this.repoRoot = options.repoRoot;
    this.modelName = options.model ?? options.config.defaultModel;
    this.mode = options.mode ?? 'code';
    this.maxIterations = options.maxIterations ?? options.config.maxIterations;
    this.tools = options.tools ?? createDefaultRegistry();
    this.onToolCall = options.onToolCall;
    this.onToolResult = options.onToolResult;
    this.onText = options.onText;
  }

  /**
   * Run the agent with a user message (non-streaming)
   */
  async run(userMessage: string): Promise<AgentResult> {
    this.messages.push({ role: 'user', content: userMessage });

    const systemPrompt = buildSystemPrompt({
      mode: this.mode,
      repoRoot: this.repoRoot,
      availableTools: this.tools.getAll(),
    });

    const model = getModel(this.modelName);
    const toolContext: ToolContext = {
      repoRoot: this.repoRoot,
      workingDir: this.repoRoot,
      logger: this.logger,
    };

    const toolCalls: ToolCallRecord[] = [];
    let totalUsage = { promptTokens: 0, completionTokens: 0 };
    let iterations = 0;

    this.logger.info({ model: this.modelName, mode: this.mode }, 'Starting agent run');

    const result = await generateText({
      model,
      system: systemPrompt,
      messages: this.messages,
      tools: this.tools.toAITools(toolContext),
      maxSteps: this.maxIterations,
      onStepFinish: (step: StepResult<Record<string, Tool>>) => {
        iterations++;

        if (step.usage) {
          totalUsage.promptTokens += step.usage.promptTokens;
          totalUsage.completionTokens += step.usage.completionTokens;
        }

        if (step.text && this.onText) {
          this.onText(step.text);
        }

        if (step.toolCalls) {
          for (const call of step.toolCalls) {
            this.logger.debug({ tool: call.toolName, args: call.args }, 'Tool call');
            this.onToolCall?.(call.toolName, call.args);
          }
        }

        if (step.toolResults) {
          for (const result of step.toolResults) {
            this.logger.debug({ tool: result.toolName }, 'Tool result');
            this.onToolResult?.(result.toolName, result.result);

            // Find matching call
            const call = step.toolCalls?.find((c) => c.toolCallId === result.toolCallId);
            if (call) {
              toolCalls.push({
                name: call.toolName,
                input: call.args,
                result: result.result,
              });
            }
          }
        }
      },
    });

    // Add assistant response to history
    this.messages.push({ role: 'assistant', content: result.text });

    this.logger.info(
      { iterations, toolCalls: toolCalls.length, usage: totalUsage },
      'Agent run complete'
    );

    return {
      response: result.text,
      toolCalls,
      usage: totalUsage,
      iterations,
    };
  }

  /**
   * Run the agent with streaming output
   */
  async *runStream(userMessage: string): AsyncGenerator<AgentStreamEvent> {
    this.messages.push({ role: 'user', content: userMessage });

    const systemPrompt = buildSystemPrompt({
      mode: this.mode,
      repoRoot: this.repoRoot,
      availableTools: this.tools.getAll(),
    });

    const model = getModel(this.modelName);
    const toolContext: ToolContext = {
      repoRoot: this.repoRoot,
      workingDir: this.repoRoot,
      logger: this.logger,
    };

    const result = streamText({
      model,
      system: systemPrompt,
      messages: this.messages,
      tools: this.tools.toAITools(toolContext),
      maxSteps: this.maxIterations,
    });

    let fullText = '';

    for await (const event of result.fullStream) {
      switch (event.type) {
        case 'text-delta':
          fullText += event.textDelta;
          yield { type: 'text-delta', textDelta: event.textDelta };
          break;

        case 'tool-call':
          yield {
            type: 'tool-call',
            toolName: event.toolName,
            args: event.args,
            toolCallId: event.toolCallId,
          };
          break;

        case 'tool-result':
          yield {
            type: 'tool-result',
            toolName: event.toolName,
            result: event.result,
            toolCallId: event.toolCallId,
          };
          break;

        case 'finish':
          yield {
            type: 'finish',
            usage: event.usage,
          };
          break;

        case 'error':
          yield {
            type: 'error',
            error: event.error,
          };
          break;
      }
    }

    // Add assistant response to history
    if (fullText) {
      this.messages.push({ role: 'assistant', content: fullText });
    }
  }

  /**
   * Get the current message history
   */
  getMessages(): CoreMessage[] {
    return [...this.messages];
  }

  /**
   * Clear the message history
   */
  clearHistory(): void {
    this.messages = [];
  }

  /**
   * Add a message to history (for resuming sessions)
   */
  addMessage(message: CoreMessage): void {
    this.messages.push(message);
  }
}

export type AgentStreamEvent =
  | { type: 'text-delta'; textDelta: string }
  | { type: 'tool-call'; toolName: string; args: unknown; toolCallId: string }
  | { type: 'tool-result'; toolName: string; result: unknown; toolCallId: string }
  | { type: 'finish'; usage: { promptTokens: number; completionTokens: number } }
  | { type: 'error'; error: unknown };
