import OpenAI from "openai"
import { ToolRegistry, type ToolContext } from "../tools/registry"
import { Conversation, type ToolCallMessage } from "./conversation"
import { buildContextPrompt } from "../context/builder"
import chalk from "chalk"

export interface AgentLoopConfig {
  model?: string
  maxTurns?: number
  verbose?: boolean
  cwd: string
  workspaceRoot: string
}

export interface AgentLoopResult {
  finalText: string
  totalToolCalls: number
  totalTokensIn: number
  totalTokensOut: number
  turns: number
}

export async function runAgentLoop(
  client: OpenAI,
  conversation: Conversation,
  registry: ToolRegistry,
  userInput: string,
  config: AgentLoopConfig,
  onTextChunk?: (text: string) => void
): Promise<AgentLoopResult> {
  const model = config.model ?? "z-ai/GLM-5.1"
  const maxTurns = config.maxTurns ?? 20

  conversation.addUserMessage(userInput)

  const toolSpecs = registry.listSpecs()
  const openaiTools: OpenAI.ChatCompletionTool[] = toolSpecs.map((spec) => ({
    type: "function" as const,
    function: {
      name: spec.name,
      description: spec.description,
      parameters: spec.inputSchema as OpenAI.FunctionParameters,
    },
  }))

  const contextResult = await buildContextPrompt(config.workspaceRoot, config.cwd)
  const contextPrompt = contextResult.prompt

  let totalToolCalls = 0
  let totalTokensIn = 0
  let totalTokensOut = 0
  let turns = 0
  let finalText = ""

  for (let turn = 0; turn < maxTurns; turn++) {
    const apiMessages = conversation.getApiMessages()
    const systemMessage: OpenAI.ChatCompletionSystemMessageParam = {
      role: "system",
      content: contextPrompt,
    }

    const allMessages: OpenAI.ChatCompletionMessageParam[] = [
      systemMessage,
      ...apiMessages.map((msg): OpenAI.ChatCompletionMessageParam => {
        if (msg.role === "system") return { role: "system", content: msg.content }
        if (msg.role === "user") return { role: "user", content: msg.content }
        if (msg.role === "assistant") {
          const m: OpenAI.ChatCompletionAssistantMessageParam = {
            role: "assistant",
            content: msg.content || null,
          }
          if (msg.toolCalls && msg.toolCalls.length > 0) {
            m.tool_calls = msg.toolCalls.map((tc): OpenAI.ChatCompletionMessageToolCall => ({
              id: tc.id,
              type: "function",
              function: { name: tc.function.name, arguments: tc.function.arguments },
            }))
          }
          return m
        }
        if (msg.role === "tool") {
          return {
            role: "tool",
            tool_call_id: msg.toolCallId ?? "",
            content: typeof msg.content === "string" ? msg.content : JSON.stringify(msg.content),
          } as OpenAI.ChatCompletionToolMessageParam
        }
        return { role: "user", content: msg.content }
      }),
    ]

    const stream = await client.chat.completions.create({
      model,
      messages: allMessages,
      tools: openaiTools.length > 0 ? openaiTools : undefined,
      stream: true,
      temperature: 0.7,
      max_tokens: 16384,
    })

    let assistantContent = ""
    const toolCalls: ToolCallMessage[] = []
    let currentToolCall: { id: string; name: string; arguments: string } | null = null

    for await (const chunk of stream) {
      const delta = chunk.choices[0]?.delta
      if (!delta) continue

      if (delta.content) {
        assistantContent += delta.content
        if (onTextChunk) onTextChunk(delta.content)
        else process.stdout.write(delta.content)
      }

      if (delta.tool_calls) {
        for (const tc of delta.tool_calls) {
          if (tc.id && tc.function?.name) {
            currentToolCall = {
              id: tc.id,
              name: tc.function.name,
              arguments: tc.function.arguments ?? "",
            }
            toolCalls.push({
              id: tc.id,
              type: "function",
              function: { name: tc.function.name, arguments: tc.function.arguments ?? "" },
            })
          } else if (currentToolCall && tc.function?.arguments) {
            currentToolCall.arguments += tc.function.arguments
            const last = toolCalls[toolCalls.length - 1]
            if (last) last.function.arguments = currentToolCall.arguments
          }
        }
      }

      if (chunk.usage) {
        totalTokensIn += chunk.usage.prompt_tokens ?? 0
        totalTokensOut += chunk.usage.completion_tokens ?? 0
      }
    }

    if (!onTextChunk) console.log()
    turns++

    if (toolCalls.length === 0) {
      conversation.addAssistantMessage(assistantContent)
      finalText = assistantContent
      break
    }

    conversation.addAssistantMessage(assistantContent, toolCalls)

    for (const tc of toolCalls) {
      totalToolCalls++
      let toolInput: Record<string, unknown>
      try {
        toolInput = JSON.parse(tc.function.arguments)
      } catch {
        toolInput = { raw: tc.function.arguments }
      }

      const summary = summarizeToolInput(tc.function.name, toolInput)
      console.log(chalk.cyan(`  ▶ ${tc.function.name}`) + chalk.dim(` ${summary}`))

      const toolCtx: ToolContext = {
        workspaceRoot: config.workspaceRoot,
        cwd: config.cwd,
      }

      const result = await registry.dispatch(tc.function.name, toolInput, toolCtx)

      if (result.isError) {
        console.log(chalk.red(`  ✗ ${tc.function.name}: ${String(result.output).slice(0, 200)}`))
      } else {
        console.log(chalk.green(`  ✓ ${summarizeToolResult(tc.function.name, result.output)}`))
      }

      const resultStr = typeof result.output === "string" ? result.output : JSON.stringify(result.output)
      conversation.addToolResult(tc.id, resultStr)
    }

    if (turn === maxTurns - 1) {
      finalText = "[Max tool turns reached]"
    }
  }

  return { finalText, totalToolCalls, totalTokensIn, totalTokensOut, turns }
}

function summarizeToolInput(name: string, input: Record<string, unknown>): string {
  const n = name.toLowerCase()
  if (n === "bash") {
    const cmd = String(input.command ?? "").replace(/\n/g, " ")
    return cmd.length > 60 ? cmd.slice(0, 57) + "..." : cmd
  }
  if (n === "read" || n === "write" || n === "edit") {
    return String(input.file_path ?? input.path ?? "")
  }
  if (n === "glob") return String(input.pattern ?? "")
  if (n === "grep") return String(input.pattern ?? "")
  return ""
}

function summarizeToolResult(name: string, output: unknown): string {
  if (typeof output === "string") {
    return output.length > 100 ? output.slice(0, 97) + "..." : output
  }
  if (typeof output === "object" && output !== null) {
    const obj = output as Record<string, unknown>
    if (obj.error) return `error: ${String(obj.error).slice(0, 80)}`
    if (obj.numFiles !== undefined) return `${name} · ${obj.numFiles} results`
    if (obj.exit_code !== undefined) return `${name} · exit ${obj.exit_code}`
  }
  return `${name} completed`
}