import { Installation } from "@/installation"
import { Provider } from "@/provider/provider"
import { Log } from "@/util/log"
import {
  stepCountIs,
  streamText,
  wrapLanguageModel,
  type ModelMessage,
  type StreamTextResult,
  type Tool,
  type ToolSet,
  tool,
  jsonSchema,
} from "ai"
import { clone, mergeDeep, pipe } from "remeda"
import { ProviderTransform } from "@/provider/transform"
import { Config } from "@/config/config"
import { Instance } from "@/project/instance"
import type { Agent } from "@/agent/agent"
import type { MessageV2 } from "./message-v2"
import { Plugin } from "@/plugin"
import { SystemPrompt } from "./system"
import { Flag } from "@/flag/flag"
import { PermissionNext } from "@/permission/next"
import { Auth } from "@/auth"
import { RateLimit } from "@/rate-limit/rate-limit"
import { ProviderSwitch } from "@/provider-switch/provider-switch"

export namespace LLM {
  const log = Log.create({ service: "llm" })

  export const OUTPUT_TOKEN_MAX = Flag.OPENCODE_EXPERIMENTAL_OUTPUT_TOKEN_MAX || 32_000

  export type StreamInput = {
    user: MessageV2.User
    sessionID: string
    model: Provider.Model
    agent: Agent.Info
    system: string[]
    abort: AbortSignal
    messages: ModelMessage[]
    small?: boolean
    tools: Record<string, Tool>
    retries?: number
    allProviders?: Record<string, Provider.Info>
  }

  export type StreamOutput = {
    result: StreamTextResult<ToolSet, unknown>
    singleStepTools: boolean
  }

  export type RetryState = {
    attempt: number
    totalWaitTime: number
    lastError: Error | null
    switchedProvider: boolean
  }

  async function streamWithRetry(
    input: StreamInput,
    config: {
      getLanguage: () => Promise<any>
      cfg: any
      provider: Provider.Info
      auth: any
      isCodex: boolean
      system: string[]
      params: any
      options: any
      headers: any
      maxOutputTokens: any
      tools: Record<string, Tool>
      rateLimitConfig: RateLimit.Config
    },
    retryState: RetryState,
  ): Promise<StreamOutput> {
    const { getLanguage, cfg, provider, auth, isCodex, system, params, options, headers, maxOutputTokens, tools, rateLimitConfig } = config
    const activeTools = Object.keys(tools).filter((x) => x !== "invalid")
    const singleStepTools = shouldLimitToolLoop({
      model: input.model,
      provider,
      options,
      activeTools,
    })

    while (retryState.attempt <= (rateLimitConfig.maxRetries || 0)) {
      try {
        log.info("Stream attempt", {
          attempt: retryState.attempt,
          maxRetries: rateLimitConfig.maxRetries,
          providerID: input.model.providerID,
          modelID: input.model.id,
        })

        const language = await getLanguage()

        const result = streamText({
          onError(error) {
            log.error("Stream error", { error, attempt: retryState.attempt })
            retryState.lastError = error
          },
          async experimental_repairToolCall(failed) {
            const lower = failed.toolCall.toolName.toLowerCase()
            if (lower !== failed.toolCall.toolName && tools[lower]) {
              log.info("repairing tool call", {
                tool: failed.toolCall.toolName,
                repaired: lower,
              })
              return {
                ...failed.toolCall,
                toolName: lower,
              }
            }
            return {
              ...failed.toolCall,
              input: JSON.stringify({
                tool: failed.toolCall.toolName,
                error: failed.error.message,
              }),
              toolName: "invalid",
            }
          },
          temperature: params.temperature,
          topP: params.topP,
          topK: params.topK,
          providerOptions: ProviderTransform.providerOptions(input.model, params.options),
          activeTools,
          tools,
          // Chutes accounts each provider round-trip as a separate request, so keep SDK orchestration to one step.
          stopWhen: singleStepTools ? stepCountIs(1) : undefined,
          maxOutputTokens,
          abortSignal: input.abort,
          headers: {
            ...(input.model.providerID.startsWith("opencode")
              ? {
                  "x-opencode-project": Instance.project.id,
                  "x-opencode-session": input.sessionID,
                  "x-opencode-request": input.user.id,
                  "x-opencode-client": Flag.OPENCODE_CLIENT,
                }
              : input.model.providerID !== "anthropic"
                ? {
                    "User-Agent": `opencode/${Installation.VERSION}`,
                  }
                : undefined),
            ...input.model.headers,
            ...headers,
          },
          maxRetries: 0,
          messages: [
            ...system.map(
              (x): ModelMessage => ({
                role: "system",
                content: x,
              }),
            ),
            ...input.messages,
          ],
          model: wrapLanguageModel({
            model: language,
            middleware: [
              {
                async transformParams(args) {
                  if (args.type === "stream") {
                    args.params.prompt = ProviderTransform.message(args.params.prompt, input.model, options)
                  }
                  return args.params
                },
              },
            ],
          }),
          experimental_telemetry: {
            isEnabled: cfg.experimental?.openTelemetry,
            metadata: {
              userId: cfg.username ?? "unknown",
              sessionId: input.sessionID,
            },
          },
        })

        ProviderSwitch.recordSuccess(input.model.providerID, input.model.id)

        return { result, singleStepTools }
      } catch (error) {
        retryState.lastError = error as Error
        ProviderSwitch.recordFailure(input.model.providerID, input.model.id)

        if (!RateLimit.isRateLimitError(error)) {
          log.error("Non-rate-limit error, not retrying", { error })
          throw error
        }

        const errorWithCause = error as any
        const headers = errorWithCause?.cause?.headers
        const parsedRateLimit = headers ? RateLimit.parseHeaders(headers) : null

        if (!parsedRateLimit) {
          log.warn("Rate limit error but could not parse headers", { error })
          
          if (retryState.attempt >= (rateLimitConfig.maxRetries || 0)) {
            throw error
          }

          const backoffMs = RateLimit.calculateBackoff(retryState.attempt, rateLimitConfig.backoffStrategy)
          retryState.totalWaitTime += backoffMs

          if (retryState.totalWaitTime > rateLimitConfig.maxWaitTime) {
            log.warn("Max wait time exceeded, giving up", {
              totalWaitTime: retryState.totalWaitTime,
              maxWaitTime: rateLimitConfig.maxWaitTime,
            })
            throw error
          }

          log.info("Waiting before retry (no rate limit headers)", {
            attempt: retryState.attempt,
            backoffMs,
            totalWaitTime: retryState.totalWaitTime,
          })

          await new Promise<void>((resolve) => setTimeout(resolve, backoffMs))
          retryState.attempt++
          continue
        }

        if (parsedRateLimit.type === "short-term") {
          if (retryState.attempt >= (rateLimitConfig.maxRetries || 0)) {
            log.warn("Max retries exceeded", {
              attempt: retryState.attempt,
              maxRetries: rateLimitConfig.maxRetries,
            })
            throw error
          }

          retryState.totalWaitTime += parsedRateLimit.waitTime

          if (retryState.totalWaitTime > rateLimitConfig.maxWaitTime) {
            log.warn("Max wait time exceeded", {
              totalWaitTime: retryState.totalWaitTime,
              maxWaitTime: rateLimitConfig.maxWaitTime,
            })
            throw error
          }

          log.info("Short-term rate limit, retrying", {
            attempt: retryState.attempt,
            waitTime: parsedRateLimit.waitTime,
            retryAt: parsedRateLimit.retryAt,
          })

          await RateLimit.wait(parsedRateLimit, input.abort)
          retryState.attempt++
        } else if (
          rateLimitConfig.enableProviderSwitch &&
          parsedRateLimit.waitTime >= rateLimitConfig.switchThreshold &&
          !retryState.switchedProvider &&
          input.allProviders
        ) {
          log.info("Long-term rate limit, switching provider", {
            waitTime: parsedRateLimit.waitTime,
            switchThreshold: rateLimitConfig.switchThreshold,
          })

          const nextProvider = await ProviderSwitch.trySwitchProvider(
            input.model.providerID,
            input.model.id,
            input.allProviders,
          )

          if (nextProvider) {
            return streamWithProvider(
              {
                ...input,
                model: await Provider.getModel(nextProvider.providerID, nextProvider.modelID),
              },
              retryState,
            )
          }

          log.warn("No fallback provider available, waiting", { providerID: input.model.providerID })

          if (parsedRateLimit.waitTime <= rateLimitConfig.maxWaitTime) {
            await RateLimit.wait(parsedRateLimit, input.abort)
          } else {
            throw error
          }
        } else {
          log.warn("Rate limit but no retry/switch possible", {
            type: parsedRateLimit.type,
            waitTime: parsedRateLimit.waitTime,
            maxWaitTime: rateLimitConfig.maxWaitTime,
            enableProviderSwitch: rateLimitConfig.enableProviderSwitch,
            switchThreshold: rateLimitConfig.switchThreshold,
          })
          throw error
        }
      }
    }

    throw retryState.lastError || new Error("Max retries exceeded")
  }

  async function streamWithProvider(
    input: StreamInput,
    retryState: RetryState = { attempt: 0, totalWaitTime: 0, lastError: null, switchedProvider: false },
  ): Promise<StreamOutput> {
    const l = log
      .clone()
      .tag("providerID", input.model.providerID)
      .tag("modelID", input.model.id)
      .tag("sessionID", input.sessionID)
      .tag("small", (input.small ?? false).toString())
      .tag("agent", input.agent.name)
      .tag("mode", input.agent.mode)

    l.info("stream", {
      modelID: input.model.id,
      providerID: input.model.providerID,
    })

    const [language, cfg, provider, auth] = await Promise.all([
      Provider.getLanguage(input.model),
      Config.get(),
      Provider.getProvider(input.model.providerID),
      Auth.get(input.model.providerID),
    ])
    const isCodex = provider.id === "openai" && auth?.type === "oauth"

    const isWordPress = await SystemPrompt.detectWordPressProject()
    const systemPrompts = await SystemPrompt.provider(input.model, isWordPress)
    
    const system = []
    system.push(
      [
        ...(input.agent.prompt ? [input.agent.prompt] : isCodex ? [] : systemPrompts),
        ...input.system,
        ...(input.user.system ? [input.user.system] : []),
      ]
        .filter((x) => x)
        .join("\n"),
    )

    const header = system[0]
    const original = clone(system)
    await Plugin.trigger(
      "experimental.chat.system.transform",
      { sessionID: input.sessionID, model: input.model },
      { system },
    )
    if (system.length === 0) {
      system.push(...original)
    }
    if (system.length > 2 && system[0] === header) {
      const rest = system.slice(1)
      system.length = 0
      system.push(header, rest.join("\n"))
    }

    const variant =
      !input.small && input.model.variants && input.user.variant ? input.model.variants[input.user.variant] : {}
    const base = input.small
      ? ProviderTransform.smallOptions(input.model)
      : ProviderTransform.options({
          model: input.model,
          sessionID: input.sessionID,
          providerOptions: provider.options,
        })
    const options: Record<string, any> = pipe(
      base,
      mergeDeep(input.model.options),
      mergeDeep(input.agent.options),
      mergeDeep(variant),
    )
    if (isCodex) {
      options.instructions = SystemPrompt.instructions()
    }

    const params = await Plugin.trigger(
      "chat.params",
      {
        sessionID: input.sessionID,
        agent: input.agent,
        model: input.model,
        provider,
        message: input.user,
      },
      {
        temperature: input.model.capabilities.temperature
          ? (input.agent.temperature ?? ProviderTransform.temperature(input.model))
          : undefined,
        topP: input.agent.topP ?? ProviderTransform.topP(input.model),
        topK: ProviderTransform.topK(input.model),
        options,
      },
    )

    const { headers: customHeaders } = await Plugin.trigger(
      "chat.headers",
      {
        sessionID: input.sessionID,
        agent: input.agent,
        model: input.model,
        provider,
        message: input.user,
      },
      {
        headers: {},
      },
    )

    const maxOutputTokens =
      isCodex || provider.id.includes("github-copilot")
        ? undefined
        : ProviderTransform.maxOutputTokens(
            input.model.api.npm,
            params.options,
            input.model.limit.output,
            OUTPUT_TOKEN_MAX,
          )

    const tools = await resolveTools(input as Pick<StreamInput, "tools" | "agent" | "user">)

    const isLiteLLMProxy =
      provider.options?.["litellmProxy"] === true ||
      input.model.providerID.toLowerCase().includes("litellm") ||
      input.model.api.id.toLowerCase().includes("litellm")

    if (isLiteLLMProxy && Object.keys(tools).length === 0 && hasToolCalls(input.messages)) {
      tools["_noop"] = tool({
        description:
          "Placeholder for LiteLLM/Anthropic proxy compatibility - required when message history contains tool calls but no active tools are needed",
        inputSchema: jsonSchema({ type: "object", properties: {} }),
        execute: async () => ({ output: "", title: "", metadata: {} }),
      })
    }

    const rateLimitConfig = RateLimit.parseConfig(cfg.rateLimit || {})

    return streamWithRetry(
      input,
      {
        getLanguage: () => Provider.getLanguage(input.model),
        cfg,
        provider,
        auth,
        isCodex,
        system,
        params,
        options,
        headers: customHeaders,
        maxOutputTokens,
        tools,
        rateLimitConfig,
      },
      retryState,
    )
  }

  export async function stream(input: StreamInput): Promise<StreamOutput> {
    return streamWithProvider(input)
  }

  export function shouldLimitToolLoopForModel(model: Provider.Model): boolean {
    if (model.providerID === "chutes") return true
    const url = model.api.url?.toLowerCase() || ""
    return url.includes("chutes.ai") || url.includes("/chutes/")
  }

  async function resolveTools(input: Pick<StreamInput, "tools" | "agent" | "user">) {
    const disabled = PermissionNext.disabled(Object.keys(input.tools), input.agent.permission)
    for (const tool of Object.keys(input.tools)) {
      if (input.user.tools?.[tool] === false || disabled.has(tool)) {
        delete input.tools[tool]
      }
    }
    return input.tools
  }

  export function hasToolCalls(messages: ModelMessage[]): boolean {
    for (const msg of messages) {
      if (!Array.isArray(msg.content)) continue
      for (const part of msg.content) {
        if (part.type === "tool-call" || part.type === "tool-result") return true
      }
    }
    return false
  }

  function shouldLimitToolLoop(input: {
    model: Provider.Model
    provider: Provider.Info
    options: Record<string, any>
    activeTools: string[]
  }) {
    if (input.activeTools.length === 0) return false
    if (input.model.providerID === "chutes") return true

    const url = [
      input.model.api.url,
      input.provider.options?.baseURL,
      input.provider.options?.baseUrl,
      input.options?.baseURL,
      input.options?.baseUrl,
    ]
      .filter((x) => typeof x === "string")
      .map((x) => x.toLowerCase())

    return url.some((x) => x.includes("chutes.ai") || x.includes("/chutes/"))
  }
}
