import { z } from "zod"
import { Hono, MiddlewareHandler } from "hono"
import { cors } from "hono/cors"
import { HTTPException } from "hono/http-exception"
import { zValidator } from "@hono/zod-validator"
import { Resource } from "sst"
import { generateText, streamText } from "ai"
import { createAnthropic } from "@ai-sdk/anthropic"
import { createOpenAI } from "@ai-sdk/openai"
import { createOpenAICompatible } from "@ai-sdk/openai-compatible"
import type { LanguageModelV2Usage, LanguageModelV2Prompt } from "@ai-sdk/provider"
import { type ChatCompletionCreateParamsBase } from "openai/resources/chat/completions"
import { Actor } from "@opencode/cloud-core/actor.js"
import { and, Database, eq, sql } from "@opencode/cloud-core/drizzle/index.js"
import { UserTable } from "@opencode/cloud-core/schema/user.sql.js"
import { KeyTable } from "@opencode/cloud-core/schema/key.sql.js"
import { createClient } from "@openauthjs/openauth/client"
import { Log } from "@opencode/cloud-core/util/log.js"
import { Billing } from "@opencode/cloud-core/billing.js"
import { Workspace } from "@opencode/cloud-core/workspace.js"
import { BillingTable, PaymentTable, UsageTable } from "@opencode/cloud-core/schema/billing.sql.js"
import { centsToMicroCents } from "@opencode/cloud-core/util/price.js"
import { Identifier } from "../../core/src/identifier"

type Env = {}

let _client: ReturnType<typeof createClient>
const client = () => {
  if (_client) return _client
  _client = createClient({
    clientID: "api",
    issuer: Resource.AUTH_API_URL.value,
  })
  return _client
}

const SUPPORTED_MODELS = {
  "anthropic/claude-sonnet-4": {
    input: 0.0000015,
    output: 0.000006,
    reasoning: 0.0000015,
    cacheRead: 0.0000001,
    cacheWrite: 0.0000001,
    model: () =>
      createAnthropic({
        apiKey: Resource.ANTHROPIC_API_KEY.value,
      })("claude-sonnet-4-20250514"),
  },
  "openai/gpt-4.1": {
    input: 0.0000015,
    output: 0.000006,
    reasoning: 0.0000015,
    cacheRead: 0.0000001,
    cacheWrite: 0.0000001,
    model: () =>
      createOpenAI({
        apiKey: Resource.OPENAI_API_KEY.value,
      })("gpt-4.1"),
  },
  "zhipuai/glm-4.5-flash": {
    input: 0,
    output: 0,
    reasoning: 0,
    cacheRead: 0,
    cacheWrite: 0,
    model: () =>
      createOpenAICompatible({
        name: "Zhipu AI",
        baseURL: "https://api.z.ai/api/paas/v4",
        apiKey: Resource.ZHIPU_API_KEY.value,
      })("glm-4.5-flash"),
  },
}

const log = Log.create({
  namespace: "api",
})

const GatewayAuth: MiddlewareHandler = async (c, next) => {
  const authHeader = c.req.header("authorization")

  if (!authHeader || !authHeader.startsWith("Bearer ")) {
    return c.json(
      {
        error: {
          message: "Missing API key.",
          type: "invalid_request_error",
          param: null,
          code: "unauthorized",
        },
      },
      401,
    )
  }

  const apiKey = authHeader.split(" ")[1]

  // Check against KeyTable
  const keyRecord = await Database.use((tx) =>
    tx
      .select({
        id: KeyTable.id,
        workspaceID: KeyTable.workspaceID,
      })
      .from(KeyTable)
      .where(eq(KeyTable.key, apiKey))
      .then((rows) => rows[0]),
  )

  if (!keyRecord) {
    return c.json(
      {
        error: {
          message: "Invalid API key.",
          type: "invalid_request_error",
          param: null,
          code: "unauthorized",
        },
      },
      401,
    )
  }

  c.set("keyRecord", keyRecord)
  await next()
}

const RestAuth: MiddlewareHandler = async (c, next) => {
  const authorization = c.req.header("authorization")
  if (!authorization) {
    return Actor.provide("public", {}, next)
  }
  const token = authorization.split(" ")[1]
  if (!token)
    throw new HTTPException(403, {
      message: "Bearer token is required.",
    })

  const verified = await client().verify(token)
  if (verified.err) {
    throw new HTTPException(403, {
      message: "Invalid token.",
    })
  }
  let subject = verified.subject as Actor.Info
  if (subject.type === "account") {
    const workspaceID = c.req.header("x-opencode-workspace")
    const email = subject.properties.email
    if (workspaceID) {
      const user = await Database.use((tx) =>
        tx
          .select({
            id: UserTable.id,
            workspaceID: UserTable.workspaceID,
            email: UserTable.email,
          })
          .from(UserTable)
          .where(and(eq(UserTable.email, email), eq(UserTable.workspaceID, workspaceID)))
          .then((rows) => rows[0]),
      )
      if (!user)
        throw new HTTPException(403, {
          message: "You do not have access to this workspace.",
        })
      subject = {
        type: "user",
        properties: {
          userID: user.id,
          workspaceID: workspaceID,
          email: user.email,
        },
      }
    }
  }
  await Actor.provide(subject.type, subject.properties, next)
}

const app = new Hono<{ Bindings: Env; Variables: { keyRecord?: { id: string; workspaceID: string } } }>()
  .get("/", (c) => c.text("Hello, world!"))
  .post("/v1/chat/completions", GatewayAuth, async (c) => {
    try {
      const body = await c.req.json<ChatCompletionCreateParamsBase>()

      console.log(body)

      const model = SUPPORTED_MODELS[body.model as keyof typeof SUPPORTED_MODELS]?.model()
      if (!model) throw new Error(`Unsupported model: ${body.model}`)

      const requestBody = transformOpenAIRequestToAiSDK()

      return body.stream ? await handleStream() : await handleGenerate()

      async function handleStream() {
        const result = await streamText({
          model,
          ...requestBody,
        })

        const encoder = new TextEncoder()
        const stream = new ReadableStream({
          async start(controller) {
            const id = `chatcmpl-${Date.now()}`
            const created = Math.floor(Date.now() / 1000)

            try {
              for await (const chunk of result.fullStream) {
                // TODO
                //console.log("!!! CHUCK !!!", chunk);
                switch (chunk.type) {
                  case "text-delta": {
                    const data = {
                      id,
                      object: "chat.completion.chunk",
                      created,
                      model: body.model,
                      choices: [
                        {
                          index: 0,
                          delta: {
                            content: chunk.text,
                          },
                          finish_reason: null,
                        },
                      ],
                    }
                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`))
                    break
                  }

                  case "reasoning-delta": {
                    const data = {
                      id,
                      object: "chat.completion.chunk",
                      created,
                      model: body.model,
                      choices: [
                        {
                          index: 0,
                          delta: {
                            reasoning_content: chunk.text,
                          },
                          finish_reason: null,
                        },
                      ],
                    }
                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`))
                    break
                  }

                  case "tool-call": {
                    const data = {
                      id,
                      object: "chat.completion.chunk",
                      created,
                      model: body.model,
                      choices: [
                        {
                          index: 0,
                          delta: {
                            tool_calls: [
                              {
                                id: chunk.toolCallId,
                                type: "function",
                                function: {
                                  name: chunk.toolName,
                                  arguments: JSON.stringify(chunk.input),
                                },
                              },
                            ],
                          },
                          finish_reason: null,
                        },
                      ],
                    }
                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`))
                    break
                  }

                  case "error": {
                    const data = {
                      id,
                      object: "chat.completion.chunk",
                      created,
                      model: body.model,
                      error: {
                        message: chunk.error,
                        type: "server_error",
                      },
                    }
                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`))
                    controller.enqueue(encoder.encode("data: [DONE]\n\n"))
                    controller.close()
                    break
                  }

                  case "finish": {
                    const finishReason =
                      {
                        stop: "stop",
                        length: "length",
                        "content-filter": "content_filter",
                        "tool-calls": "tool_calls",
                        error: "stop",
                        other: "stop",
                        unknown: "stop",
                      }[chunk.finishReason] || "stop"

                    const data = {
                      id,
                      object: "chat.completion.chunk",
                      created,
                      model: body.model,
                      choices: [
                        {
                          index: 0,
                          delta: {},
                          finish_reason: finishReason,
                        },
                      ],
                      usage: {
                        prompt_tokens: chunk.totalUsage.inputTokens,
                        completion_tokens: chunk.totalUsage.outputTokens,
                        total_tokens: chunk.totalUsage.totalTokens,
                        completion_tokens_details: {
                          reasoning_tokens: chunk.totalUsage.reasoningTokens,
                        },
                        prompt_tokens_details: {
                          cached_tokens: chunk.totalUsage.cachedInputTokens,
                        },
                      },
                    }
                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(data)}\n\n`))
                    controller.enqueue(encoder.encode("data: [DONE]\n\n"))
                    controller.close()
                    break
                  }

                  //case "stream-start":
                  //case "response-metadata":
                  case "start-step":
                  case "finish-step":
                  case "text-start":
                  case "text-end":
                  case "reasoning-start":
                  case "reasoning-end":
                  case "tool-input-start":
                  case "tool-input-delta":
                  case "tool-input-end":
                  case "raw":
                  default:
                    // Log unknown chunk types for debugging
                    console.warn(`Unknown chunk type: ${(chunk as any).type}`)
                    break
                }
              }
            } catch (error) {
              controller.error(error)
            }
          },
        })

        return new Response(stream, {
          headers: {
            "Content-Type": "text/plain; charset=utf-8",
            "Cache-Control": "no-cache",
            Connection: "keep-alive",
          },
        })
      }

      async function handleGenerate() {
        const response = await generateText({
          model,
          ...requestBody,
        })
        await trackUsage(body.model, response.usage)
        return c.json({
          id: `chatcmpl-${Date.now()}`,
          object: "chat.completion" as const,
          created: Math.floor(Date.now() / 1000),
          model: body.model,
          choices: [
            {
              index: 0,
              message: {
                role: "assistant" as const,
                content: response.content?.find((c) => c.type === "text")?.text ?? "",
                reasoning_content: response.content?.find((c) => c.type === "reasoning")?.text,
                tool_calls: response.content
                  ?.filter((c) => c.type === "tool-call")
                  .map((toolCall) => ({
                    id: toolCall.toolCallId,
                    type: "function" as const,
                    function: {
                      name: toolCall.toolName,
                      arguments: toolCall.input,
                    },
                  })),
              },
              finish_reason:
                (
                  {
                    stop: "stop",
                    length: "length",
                    "content-filter": "content_filter",
                    "tool-calls": "tool_calls",
                    error: "stop",
                    other: "stop",
                    unknown: "stop",
                  } as const
                )[response.finishReason] || "stop",
            },
          ],
          usage: {
            prompt_tokens: response.usage?.inputTokens,
            completion_tokens: response.usage?.outputTokens,
            total_tokens: response.usage?.totalTokens,
            completion_tokens_details: {
              reasoning_tokens: response.usage?.reasoningTokens,
            },
            prompt_tokens_details: {
              cached_tokens: response.usage?.cachedInputTokens,
            },
          },
        })
      }

      function transformOpenAIRequestToAiSDK() {
        const prompt = transformMessages()

        return {
          prompt,
          maxOutputTokens: body.max_tokens ?? body.max_completion_tokens ?? undefined,
          temperature: body.temperature ?? undefined,
          topP: body.top_p ?? undefined,
          frequencyPenalty: body.frequency_penalty ?? undefined,
          presencePenalty: body.presence_penalty ?? undefined,
          providerOptions: body.reasoning_effort
            ? {
                anthropic: {
                  reasoningEffort: body.reasoning_effort,
                },
              }
            : undefined,
          stopSequences: (typeof body.stop === "string" ? [body.stop] : body.stop) ?? undefined,
          responseFormat: (() => {
            if (!body.response_format) return { type: "text" }
            if (body.response_format.type === "json_schema")
              return {
                type: "json",
                schema: body.response_format.json_schema.schema,
                name: body.response_format.json_schema.name,
                description: body.response_format.json_schema.description,
              }
            if (body.response_format.type === "json_object") return { type: "json" }
            throw new Error("Unsupported response format")
          })(),
          seed: body.seed ?? undefined,
        }

        function transformTools() {
          const { tools, tool_choice } = body

          if (!tools || tools.length === 0) {
            return { tools: undefined, toolChoice: undefined }
          }

          const aiSdkTools = tools.reduce(
            (acc, tool) => {
              acc[tool.function.name] = {
                type: "function" as const,
                name: tool.function.name,
                description: tool.function.description,
                inputSchema: tool.function.parameters,
              }
              return acc
            },
            {} as Record<string, any>,
          )

          let aiSdkToolChoice
          if (tool_choice == null) {
            aiSdkToolChoice = undefined
          } else if (tool_choice === "auto") {
            aiSdkToolChoice = "auto"
          } else if (tool_choice === "none") {
            aiSdkToolChoice = "none"
          } else if (tool_choice === "required") {
            aiSdkToolChoice = "required"
          } else if (tool_choice.type === "function") {
            aiSdkToolChoice = {
              type: "tool",
              toolName: tool_choice.function.name,
            }
          }

          return { tools: aiSdkTools, toolChoice: aiSdkToolChoice }
        }

        function transformMessages() {
          const { messages } = body
          const prompt: LanguageModelV2Prompt = []

          for (const message of messages) {
            switch (message.role) {
              case "system": {
                prompt.push({
                  role: "system",
                  content: message.content as string,
                })
                break
              }

              case "user": {
                if (typeof message.content === "string") {
                  prompt.push({
                    role: "user",
                    content: [{ type: "text", text: message.content }],
                  })
                } else {
                  const content = message.content.map((part) => {
                    switch (part.type) {
                      case "text":
                        return { type: "text" as const, text: part.text }
                      case "image_url":
                        return {
                          type: "file" as const,
                          mediaType: "image/jpeg" as const,
                          data: part.image_url.url,
                        }
                      default:
                        throw new Error(`Unsupported content part type: ${(part as any).type}`)
                    }
                  })
                  prompt.push({
                    role: "user",
                    content,
                  })
                }
                break
              }

              case "assistant": {
                const content: Array<
                  | { type: "text"; text: string }
                  | {
                      type: "tool-call"
                      toolCallId: string
                      toolName: string
                      input: any
                    }
                > = []

                if (message.content) {
                  content.push({
                    type: "text",
                    text: message.content as string,
                  })
                }

                if (message.tool_calls) {
                  for (const toolCall of message.tool_calls) {
                    content.push({
                      type: "tool-call",
                      toolCallId: toolCall.id,
                      toolName: toolCall.function.name,
                      input: JSON.parse(toolCall.function.arguments),
                    })
                  }
                }

                prompt.push({
                  role: "assistant",
                  content,
                })
                break
              }

              case "tool": {
                prompt.push({
                  role: "tool",
                  content: [
                    {
                      type: "tool-result",
                      toolName: "placeholder",
                      toolCallId: message.tool_call_id,
                      output: {
                        type: "text",
                        value: message.content as string,
                      },
                    },
                  ],
                })
                break
              }

              default: {
                throw new Error(`Unsupported message role: ${message.role}`)
              }
            }
          }

          return prompt
        }
      }

      async function trackUsage(model: string, usage: LanguageModelV2Usage) {
        const keyRecord = c.get("keyRecord")
        if (!keyRecord) return

        const modelData = SUPPORTED_MODELS[model as keyof typeof SUPPORTED_MODELS]
        if (!modelData) throw new Error(`Unsupported model: ${model}`)

        const inputCost = modelData.input * (usage.inputTokens ?? 0)
        const outputCost = modelData.output * (usage.outputTokens ?? 0)
        const reasoningCost = modelData.reasoning * (usage.reasoningTokens ?? 0)
        const cacheReadCost = modelData.cacheRead * (usage.cachedInputTokens ?? 0)
        const cacheWriteCost = modelData.cacheWrite * (usage.outputTokens ?? 0)

        const totalCost = inputCost + outputCost + reasoningCost + cacheReadCost + cacheWriteCost

        await Actor.provide("system", { workspaceID: keyRecord.workspaceID }, async () => {
          await Billing.consume({
            model,
            inputTokens: usage.inputTokens ?? 0,
            outputTokens: usage.outputTokens ?? 0,
            reasoningTokens: usage.reasoningTokens ?? 0,
            cacheReadTokens: usage.cachedInputTokens ?? 0,
            cacheWriteTokens: usage.outputTokens ?? 0,
            costInCents: totalCost * 100,
          })
        })

        await Database.use((tx) =>
          tx
            .update(KeyTable)
            .set({ timeUsed: sql`now()` })
            .where(eq(KeyTable.id, keyRecord.id)),
        )
      }
    } catch (error: any) {
      return c.json({ error: { message: error.message } }, 500)
    }
  })
  .use("/*", cors())
  .use(RestAuth)
  .get("/rest/account", async (c) => {
    const account = Actor.assert("account")
    let workspaces = await Workspace.list()
    if (workspaces.length === 0) {
      await Workspace.create()
      workspaces = await Workspace.list()
    }
    return c.json({
      id: account.properties.accountID,
      email: account.properties.email,
      workspaces,
    })
  })
  .get("/billing/info", async (c) => {
    const billing = await Billing.get()
    const payments = await Database.use((tx) =>
      tx
        .select()
        .from(PaymentTable)
        .where(eq(PaymentTable.workspaceID, Actor.workspace()))
        .orderBy(sql`${PaymentTable.timeCreated} DESC`)
        .limit(100),
    )
    const usage = await Database.use((tx) =>
      tx
        .select()
        .from(UsageTable)
        .where(eq(UsageTable.workspaceID, Actor.workspace()))
        .orderBy(sql`${UsageTable.timeCreated} DESC`)
        .limit(100),
    )
    return c.json({ billing, payments, usage })
  })
  .post(
    "/billing/checkout",
    zValidator(
      "json",
      z.custom<{
        success_url: string
        cancel_url: string
      }>(),
    ),
    async (c) => {
      const account = Actor.assert("user")

      const body = await c.req.json()

      const customer = await Billing.get()
      const session = await Billing.stripe().checkout.sessions.create({
        mode: "payment",
        line_items: [
          {
            price_data: {
              currency: "usd",
              product_data: {
                name: "OpenControl credits",
              },
              unit_amount: 2000, // $20 minimum
            },
            quantity: 1,
          },
        ],
        payment_intent_data: {
          setup_future_usage: "on_session",
        },
        ...(customer.customerID
          ? { customer: customer.customerID }
          : {
              customer_email: account.properties.email,
              customer_creation: "always",
            }),
        metadata: {
          workspaceID: Actor.workspace(),
        },
        currency: "usd",
        payment_method_types: ["card"],
        success_url: body.success_url,
        cancel_url: body.cancel_url,
      })

      return c.json({
        url: session.url,
      })
    },
  )
  .post("/billing/portal", async (c) => {
    const body = await c.req.json()

    const customer = await Billing.get()
    if (!customer?.customerID) {
      throw new Error("No stripe customer ID")
    }

    const session = await Billing.stripe().billingPortal.sessions.create({
      customer: customer.customerID,
      return_url: body.return_url,
    })

    return c.json({
      url: session.url,
    })
  })
  .post("/stripe/webhook", async (c) => {
    const body = await Billing.stripe().webhooks.constructEventAsync(
      await c.req.text(),
      c.req.header("stripe-signature")!,
      Resource.STRIPE_WEBHOOK_SECRET.value,
    )

    console.log(body.type, JSON.stringify(body, null, 2))
    if (body.type === "checkout.session.completed") {
      const workspaceID = body.data.object.metadata?.workspaceID
      const customerID = body.data.object.customer as string
      const paymentID = body.data.object.payment_intent as string
      const amount = body.data.object.amount_total

      if (!workspaceID) throw new Error("Workspace ID not found")
      if (!customerID) throw new Error("Customer ID not found")
      if (!amount) throw new Error("Amount not found")
      if (!paymentID) throw new Error("Payment ID not found")

      await Actor.provide("system", { workspaceID }, async () => {
        const customer = await Billing.get()
        if (customer?.customerID && customer.customerID !== customerID) throw new Error("Customer ID mismatch")

        // set customer metadata
        if (!customer?.customerID) {
          await Billing.stripe().customers.update(customerID, {
            metadata: {
              workspaceID,
            },
          })
        }

        // get payment method for the payment intent
        const paymentIntent = await Billing.stripe().paymentIntents.retrieve(paymentID, {
          expand: ["payment_method"],
        })
        const paymentMethod = paymentIntent.payment_method
        if (!paymentMethod || typeof paymentMethod === "string") throw new Error("Payment method not expanded")

        await Database.transaction(async (tx) => {
          await tx
            .update(BillingTable)
            .set({
              balance: sql`${BillingTable.balance} + ${centsToMicroCents(amount)}`,
              customerID,
              paymentMethodID: paymentMethod.id,
              paymentMethodLast4: paymentMethod.card!.last4,
            })
            .where(eq(BillingTable.workspaceID, workspaceID))
          await tx.insert(PaymentTable).values({
            workspaceID,
            id: Identifier.create("payment"),
            amount: centsToMicroCents(amount),
            paymentID,
            customerID,
          })
        })
      })
    }

    console.log("finished handling")

    return c.json("ok", 200)
  })
  .get("/keys", async (c) => {
    const user = Actor.assert("user")

    const keys = await Database.use((tx) =>
      tx
        .select({
          id: KeyTable.id,
          name: KeyTable.name,
          key: KeyTable.key,
          userID: KeyTable.userID,
          timeCreated: KeyTable.timeCreated,
          timeUsed: KeyTable.timeUsed,
        })
        .from(KeyTable)
        .where(eq(KeyTable.workspaceID, user.properties.workspaceID))
        .orderBy(sql`${KeyTable.timeCreated} DESC`),
    )

    return c.json({ keys })
  })
  .post("/keys", zValidator("json", z.object({ name: z.string().min(1).max(255) })), async (c) => {
    const user = Actor.assert("user")
    const { name } = c.req.valid("json")

    // Generate secret key: sk- + 64 random characters (upper, lower, numbers)
    const chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    let randomPart = ""
    for (let i = 0; i < 64; i++) {
      randomPart += chars.charAt(Math.floor(Math.random() * chars.length))
    }
    const secretKey = `sk-${randomPart}`

    const keyRecord = await Database.use((tx) =>
      tx
        .insert(KeyTable)
        .values({
          id: Identifier.create("key"),
          workspaceID: user.properties.workspaceID,
          userID: user.properties.userID,
          name,
          key: secretKey,
          timeUsed: null,
        })
        .returning(),
    )

    return c.json({
      key: secretKey,
      id: keyRecord[0].id,
      name: keyRecord[0].name,
      created: keyRecord[0].timeCreated,
    })
  })
  .delete("/keys/:id", async (c) => {
    const user = Actor.assert("user")
    const keyId = c.req.param("id")

    const result = await Database.use((tx) =>
      tx
        .delete(KeyTable)
        .where(and(eq(KeyTable.id, keyId), eq(KeyTable.workspaceID, user.properties.workspaceID)))
        .returning({ id: KeyTable.id }),
    )

    if (result.length === 0) {
      return c.json({ error: "Key not found" }, 404)
    }

    return c.json({ success: true, id: result[0].id })
  })
  .all("*", (c) => c.text("Not Found"))

export type ApiType = typeof app

export default app