import { NodeHttpServer, NodeHttpServerRequest } from "@effect/platform-node"
import * as Http from "node:http"
import { Deferred, Effect, Layer, ServiceMap, Stream } from "effect"
import * as HttpServer from "effect/unstable/http/HttpServer"
import { HttpRouter, HttpServerRequest, HttpServerResponse } from "effect/unstable/http"

export type Usage = { input: number; output: number }

type Line = Record<string, unknown>

type Hit = {
  url: URL
  body: Record<string, unknown>
}

type Wait = {
  count: number
  ready: Deferred.Deferred<void>
}

type Sse = {
  type: "sse"
  head: unknown[]
  tail: unknown[]
  wait?: PromiseLike<unknown>
  hang?: boolean
  error?: unknown
  reset?: boolean
}

type HttpError = {
  type: "http-error"
  status: number
  body: unknown
}

export type Item = Sse | HttpError

const done = Symbol("done")

function line(input: unknown) {
  if (input === done) return "data: [DONE]\n\n"
  return `data: ${JSON.stringify(input)}\n\n`
}

function tokens(input?: Usage) {
  if (!input) return
  return {
    prompt_tokens: input.input,
    completion_tokens: input.output,
    total_tokens: input.input + input.output,
  }
}

function chunk(input: { delta?: Record<string, unknown>; finish?: string; usage?: Usage }) {
  return {
    id: "chatcmpl-test",
    object: "chat.completion.chunk",
    choices: [
      {
        delta: input.delta ?? {},
        ...(input.finish ? { finish_reason: input.finish } : {}),
      },
    ],
    ...(input.usage ? { usage: tokens(input.usage) } : {}),
  } satisfies Line
}

function role() {
  return chunk({ delta: { role: "assistant" } })
}

function textLine(value: string) {
  return chunk({ delta: { content: value } })
}

function reasonLine(value: string) {
  return chunk({ delta: { reasoning_content: value } })
}

function finishLine(reason: string, usage?: Usage) {
  return chunk({ finish: reason, usage })
}

function toolStartLine(id: string, name: string) {
  return chunk({
    delta: {
      tool_calls: [
        {
          index: 0,
          id,
          type: "function",
          function: {
            name,
            arguments: "",
          },
        },
      ],
    },
  })
}

function toolArgsLine(value: string) {
  return chunk({
    delta: {
      tool_calls: [
        {
          index: 0,
          function: {
            arguments: value,
          },
        },
      ],
    },
  })
}

function bytes(input: Iterable<unknown>) {
  return Stream.fromIterable([...input].map(line)).pipe(Stream.encodeText)
}

function created(model: string) {
  return {
    type: "response.created",
    sequence_number: 1,
    response: {
      id: "resp_test",
      created_at: Math.floor(Date.now() / 1000),
      model,
      service_tier: null,
    },
  }
}

function completed(input: { seq: number; usage?: Usage }) {
  return {
    type: "response.completed",
    sequence_number: input.seq,
    response: {
      incomplete_details: null,
      service_tier: null,
      usage: {
        input_tokens: input.usage?.input ?? 0,
        input_tokens_details: { cached_tokens: null },
        output_tokens: input.usage?.output ?? 0,
        output_tokens_details: { reasoning_tokens: null },
      },
    },
  }
}

function responses(item: Sse, model: string) {
  let seq = 1
  let msg: string | undefined
  let reason: string | undefined
  let call:
    | {
        id: string
        item: string
        name: string
        args: string
      }
    | undefined
  let usage: Usage | undefined
  const lines: unknown[] = [created(model)]

  const all = [...item.head, ...item.tail]
  for (const part of all) {
    if (!part || typeof part !== "object") continue
    if (!("choices" in part) || !Array.isArray(part.choices)) continue
    const choice = part.choices[0]
    if (!choice || typeof choice !== "object") continue
    const delta = "delta" in choice && choice.delta && typeof choice.delta === "object" ? choice.delta : undefined

    if (delta && "content" in delta && typeof delta.content === "string") {
      msg ||= "msg_1"
      if (
        !lines.some(
          (item) =>
            typeof item === "object" &&
            item &&
            "type" in item &&
            item.type === "response.output_item.added" &&
            "item" in item &&
            item.item &&
            typeof item.item === "object" &&
            "id" in item.item &&
            item.item.id === msg,
        )
      ) {
        seq += 1
        lines.push({
          type: "response.output_item.added",
          sequence_number: seq,
          output_index: 0,
          item: { type: "message", id: msg },
        })
      }
      seq += 1
      lines.push({
        type: "response.output_text.delta",
        sequence_number: seq,
        item_id: msg,
        delta: delta.content,
        logprobs: null,
      })
    }

    if (delta && "reasoning_content" in delta && typeof delta.reasoning_content === "string") {
      reason ||= "rs_1"
      if (
        !lines.some(
          (item) =>
            typeof item === "object" &&
            item &&
            "type" in item &&
            item.type === "response.output_item.added" &&
            "item" in item &&
            item.item &&
            typeof item.item === "object" &&
            "id" in item.item &&
            item.item.id === reason,
        )
      ) {
        seq += 1
        lines.push({
          type: "response.output_item.added",
          sequence_number: seq,
          output_index: 0,
          item: { type: "reasoning", id: reason, encrypted_content: null },
        })
        seq += 1
        lines.push({
          type: "response.reasoning_summary_part.added",
          sequence_number: seq,
          item_id: reason,
          summary_index: 0,
        })
      }
      seq += 1
      lines.push({
        type: "response.reasoning_summary_text.delta",
        sequence_number: seq,
        item_id: reason,
        summary_index: 0,
        delta: delta.reasoning_content,
      })
    }

    if (delta && "tool_calls" in delta && Array.isArray(delta.tool_calls)) {
      for (const tool of delta.tool_calls) {
        if (!tool || typeof tool !== "object") continue
        const fn = "function" in tool && tool.function && typeof tool.function === "object" ? tool.function : undefined
        const id = "id" in tool && typeof tool.id === "string" ? tool.id : call?.id
        const name = fn && "name" in fn && typeof fn.name === "string" ? fn.name : call?.name
        const args = fn && "arguments" in fn && typeof fn.arguments === "string" ? fn.arguments : ""
        if (!id || !name) continue
        if (!call) {
          call = { id, item: "fc_1", name, args: "" }
          seq += 1
          lines.push({
            type: "response.output_item.added",
            sequence_number: seq,
            output_index: 0,
            item: {
              type: "function_call",
              id: call.item,
              call_id: id,
              name,
              arguments: "",
              status: "in_progress",
            },
          })
        }
        call.args += args
        if (args) {
          seq += 1
          lines.push({
            type: "response.function_call_arguments.delta",
            sequence_number: seq,
            output_index: 0,
            item_id: call.item,
            delta: args,
          })
        }
      }
    }

    if ("usage" in part && part.usage && typeof part.usage === "object") {
      const raw = part.usage as Record<string, unknown>
      if (typeof raw.prompt_tokens === "number" && typeof raw.completion_tokens === "number") {
        usage = { input: raw.prompt_tokens, output: raw.completion_tokens }
      }
    }
  }

  if (msg) {
    seq += 1
    lines.push({
      type: "response.output_item.done",
      sequence_number: seq,
      output_index: 0,
      item: { type: "message", id: msg },
    })
  }
  if (reason) {
    seq += 1
    lines.push({
      type: "response.output_item.done",
      sequence_number: seq,
      output_index: 0,
      item: { type: "reasoning", id: reason, encrypted_content: null },
    })
  }
  if (call && !item.hang && !item.error) {
    seq += 1
    lines.push({
      type: "response.output_item.done",
      sequence_number: seq,
      output_index: 0,
      item: {
        type: "function_call",
        id: call.item,
        call_id: call.id,
        name: call.name,
        arguments: call.args,
        status: "completed",
      },
    })
  }
  if (!item.hang && !item.error) lines.push(completed({ seq: seq + 1, usage }))
  return { ...item, head: lines, tail: [] } satisfies Sse
}

function modelFrom(body: unknown) {
  if (!body || typeof body !== "object") return "test-model"
  if (!("model" in body) || typeof body.model !== "string") return "test-model"
  return body.model
}

function send(item: Sse) {
  const head = bytes(item.head)
  const tail = bytes([...item.tail, ...(item.hang || item.error ? [] : [done])])
  const empty = Stream.fromIterable<Uint8Array>([])
  const wait = item.wait
  const body: Stream.Stream<Uint8Array, unknown> = wait
    ? Stream.concat(head, Stream.fromEffect(Effect.promise(() => wait)).pipe(Stream.flatMap(() => tail)))
    : Stream.concat(head, tail)
  let end: Stream.Stream<Uint8Array, unknown> = empty
  if (item.error) end = Stream.concat(empty, Stream.fail(item.error))
  else if (item.hang) end = Stream.concat(empty, Stream.never)

  return HttpServerResponse.stream(Stream.concat(body, end), { contentType: "text/event-stream" })
}

const reset = Effect.fn("TestLLMServer.reset")(function* (item: Sse) {
  const req = yield* HttpServerRequest.HttpServerRequest
  const res = NodeHttpServerRequest.toServerResponse(req)
  yield* Effect.sync(() => {
    res.writeHead(200, { "content-type": "text/event-stream" })
    for (const part of item.head) res.write(line(part))
    for (const part of item.tail) res.write(line(part))
    res.destroy(new Error("connection reset"))
  })
  yield* Effect.never
})

function fail(item: HttpError) {
  return HttpServerResponse.text(JSON.stringify(item.body), {
    status: item.status,
    contentType: "application/json",
  })
}

export class Reply {
  #head: unknown[] = [role()]
  #tail: unknown[] = []
  #usage: Usage | undefined
  #finish: string | undefined
  #wait: PromiseLike<unknown> | undefined
  #hang = false
  #error: unknown
  #reset = false
  #seq = 0

  #id() {
    this.#seq += 1
    return `call_${this.#seq}`
  }

  text(value: string) {
    this.#tail = [...this.#tail, textLine(value)]
    return this
  }

  reason(value: string) {
    this.#tail = [...this.#tail, reasonLine(value)]
    return this
  }

  usage(value: Usage) {
    this.#usage = value
    return this
  }

  wait(value: PromiseLike<unknown>) {
    this.#wait = value
    return this
  }

  stop() {
    this.#finish = "stop"
    this.#hang = false
    this.#error = undefined
    this.#reset = false
    return this
  }

  toolCalls() {
    this.#finish = "tool_calls"
    this.#hang = false
    this.#error = undefined
    this.#reset = false
    return this
  }

  tool(name: string, input: unknown) {
    const id = this.#id()
    const args = JSON.stringify(input)
    this.#tail = [...this.#tail, toolStartLine(id, name), toolArgsLine(args)]
    return this.toolCalls()
  }

  pendingTool(name: string, input: unknown) {
    const id = this.#id()
    const args = JSON.stringify(input)
    const size = Math.max(1, Math.floor(args.length / 2))
    this.#tail = [...this.#tail, toolStartLine(id, name), toolArgsLine(args.slice(0, size))]
    return this
  }

  hang() {
    this.#finish = undefined
    this.#hang = true
    this.#error = undefined
    this.#reset = false
    return this
  }

  streamError(error: unknown = "boom") {
    this.#finish = undefined
    this.#hang = false
    this.#error = error
    this.#reset = false
    return this
  }

  reset() {
    this.#finish = undefined
    this.#hang = false
    this.#error = undefined
    this.#reset = true
    return this
  }

  item(): Item {
    return {
      type: "sse",
      head: this.#head,
      tail: this.#finish ? [...this.#tail, finishLine(this.#finish, this.#usage)] : this.#tail,
      wait: this.#wait,
      hang: this.#hang,
      error: this.#error,
      reset: this.#reset,
    }
  }
}

export function reply() {
  return new Reply()
}

export function httpError(status: number, body: unknown): Item {
  return {
    type: "http-error",
    status,
    body,
  }
}

export function raw(input: {
  chunks?: unknown[]
  head?: unknown[]
  tail?: unknown[]
  wait?: PromiseLike<unknown>
  hang?: boolean
  error?: unknown
  reset?: boolean
}): Item {
  return {
    type: "sse",
    head: input.head ?? input.chunks ?? [],
    tail: input.tail ?? [],
    wait: input.wait,
    hang: input.hang,
    error: input.error,
    reset: input.reset,
  }
}

function item(input: Item | Reply) {
  return input instanceof Reply ? input.item() : input
}

namespace TestLLMServer {
  export interface Service {
    readonly url: string
    readonly push: (...input: (Item | Reply)[]) => Effect.Effect<void>
    readonly text: (value: string, opts?: { usage?: Usage }) => Effect.Effect<void>
    readonly tool: (name: string, input: unknown) => Effect.Effect<void>
    readonly toolHang: (name: string, input: unknown) => Effect.Effect<void>
    readonly reason: (value: string, opts?: { text?: string; usage?: Usage }) => Effect.Effect<void>
    readonly fail: (message?: unknown) => Effect.Effect<void>
    readonly error: (status: number, body: unknown) => Effect.Effect<void>
    readonly hang: Effect.Effect<void>
    readonly hold: (value: string, wait: PromiseLike<unknown>) => Effect.Effect<void>
    readonly hits: Effect.Effect<Hit[]>
    readonly calls: Effect.Effect<number>
    readonly wait: (count: number) => Effect.Effect<void>
    readonly inputs: Effect.Effect<Record<string, unknown>[]>
    readonly pending: Effect.Effect<number>
  }
}

export class TestLLMServer extends ServiceMap.Service<TestLLMServer, TestLLMServer.Service>()("@test/LLMServer") {
  static readonly layer = Layer.effect(
    TestLLMServer,
    Effect.gen(function* () {
      const server = yield* HttpServer.HttpServer
      const router = yield* HttpRouter.HttpRouter

      let hits: Hit[] = []
      let list: Item[] = []
      let waits: Wait[] = []

      const queue = (...input: (Item | Reply)[]) => {
        list = [...list, ...input.map(item)]
      }

      const notify = Effect.fnUntraced(function* () {
        const ready = waits.filter((item) => hits.length >= item.count)
        if (!ready.length) return
        waits = waits.filter((item) => hits.length < item.count)
        yield* Effect.forEach(ready, (item) => Deferred.succeed(item.ready, void 0))
      })

      const pull = () => {
        const first = list[0]
        if (!first) return
        list = list.slice(1)
        return first
      }

      yield* router.add(
        "POST",
        "/v1/chat/completions",
        Effect.gen(function* () {
          const req = yield* HttpServerRequest.HttpServerRequest
          const next = pull()
          if (!next) return HttpServerResponse.text("unexpected request", { status: 500 })
          const body = yield* req.json.pipe(Effect.orElseSucceed(() => ({})))
          hits = [
            ...hits,
            {
              url: new URL(req.originalUrl, "http://localhost"),
              body: body && typeof body === "object" ? (body as Record<string, unknown>) : {},
            },
          ]
          yield* notify()
          if (req.originalUrl.endsWith("/v1/responses") && next.type === "sse") {
            return send(responses(next, modelFrom(body)))
          }
          if (next.type === "sse" && next.reset) {
            yield* reset(next)
            return HttpServerResponse.empty()
          }
          if (next.type === "sse") return send(next)
          return fail(next)
        }),
      )

      yield* router.add(
        "POST",
        "/v1/responses",
        Effect.gen(function* () {
          const req = yield* HttpServerRequest.HttpServerRequest
          const next = pull()
          if (!next) return HttpServerResponse.text("unexpected request", { status: 500 })
          const body = yield* req.json.pipe(Effect.orElseSucceed(() => ({})))
          hits = [
            ...hits,
            {
              url: new URL(req.originalUrl, "http://localhost"),
              body: body && typeof body === "object" ? (body as Record<string, unknown>) : {},
            },
          ]
          yield* notify()
          if (next.type === "sse") return send(responses(next, modelFrom(body)))
          return fail(next)
        }),
      )

      yield* server.serve(router.asHttpEffect())

      return TestLLMServer.of({
        url:
          server.address._tag === "TcpAddress"
            ? `http://127.0.0.1:${server.address.port}/v1`
            : `unix://${server.address.path}/v1`,
        push: Effect.fn("TestLLMServer.push")(function* (...input: (Item | Reply)[]) {
          queue(...input)
        }),
        text: Effect.fn("TestLLMServer.text")(function* (value: string, opts?: { usage?: Usage }) {
          const out = reply().text(value)
          if (opts?.usage) out.usage(opts.usage)
          queue(out.stop().item())
        }),
        tool: Effect.fn("TestLLMServer.tool")(function* (name: string, input: unknown) {
          queue(reply().tool(name, input).item())
        }),
        toolHang: Effect.fn("TestLLMServer.toolHang")(function* (name: string, input: unknown) {
          queue(reply().pendingTool(name, input).hang().item())
        }),
        reason: Effect.fn("TestLLMServer.reason")(function* (value: string, opts?: { text?: string; usage?: Usage }) {
          const out = reply().reason(value)
          if (opts?.text) out.text(opts.text)
          if (opts?.usage) out.usage(opts.usage)
          queue(out.stop().item())
        }),
        fail: Effect.fn("TestLLMServer.fail")(function* (message: unknown = "boom") {
          queue(reply().streamError(message).item())
        }),
        error: Effect.fn("TestLLMServer.error")(function* (status: number, body: unknown) {
          queue(httpError(status, body))
        }),
        hang: Effect.gen(function* () {
          queue(reply().hang().item())
        }).pipe(Effect.withSpan("TestLLMServer.hang")),
        hold: Effect.fn("TestLLMServer.hold")(function* (value: string, wait: PromiseLike<unknown>) {
          queue(reply().wait(wait).text(value).stop().item())
        }),
        hits: Effect.sync(() => [...hits]),
        calls: Effect.sync(() => hits.length),
        wait: Effect.fn("TestLLMServer.wait")(function* (count: number) {
          if (hits.length >= count) return
          const ready = yield* Deferred.make<void>()
          waits = [...waits, { count, ready }]
          yield* Deferred.await(ready)
        }),
        inputs: Effect.sync(() => hits.map((hit) => hit.body)),
        pending: Effect.sync(() => list.length),
      })
    }),
  ).pipe(Layer.provide(HttpRouter.layer), Layer.provide(NodeHttpServer.layer(() => Http.createServer(), { port: 0 })))
}