import { randomUUID } from "node:crypto";
import * as zlib from "node:zlib";
import {
  type Api,
  type AssistantMessage,
  type AssistantMessageEventStream,
  type Context,
  type Model,
  type SimpleStreamOptions,
  calculateCost,
  createAssistantMessageEventStream,
} from "@earendil-works/pi-ai";
import { clampMaxTokensToContext } from "./context-estimate.js";
import { mapContextToChat, type ChatHistoryItem, type ToolDef } from "./context-map.js";
import { anySignal, getCachedUserJwt } from "./jwt.js";
import { buildMetadata } from "./metadata.js";
import { resolveModelUid } from "./models.js";
import {
  encodeFixed64Field,
  encodeMessage,
  encodeString,
  encodeVarintField,
  frameConnectStream,
  iterFields,
} from "./wire.js";

const SOURCE_BY_ROLE: Record<string, number> = {
  user: 1,
  assistant: 2,
  tool: 4,
};

/** Connect frame payloads larger than this indicate a corrupt length prefix. */
const MAX_CONNECT_FRAME_PAYLOAD = 16 * 1024 * 1024;
/**
 * Recovery heuristic for opaque Devin `invalid_argument`/`internal error`
 * trailers: compactable message history this large is likely the cause, so the
 * failure is reported in a shape Pi recognizes as context overflow.
 */
const LARGE_HISTORY_RECOVERY_BYTES = 512 * 1024;

export type CloudChatEvent =
  | { kind: "text"; text: string }
  | { kind: "reasoning"; text: string }
  | { kind: "signature"; text: string }
  | { kind: "thinking_redacted" }
  | { kind: "message_id"; id: string }
  | { kind: "actual_model"; uid: string }
  | { kind: "tool_call_start"; id: string; name: string }
  | { kind: "tool_call_args"; argsDelta: string; id?: string }
  | { kind: "finish"; reason: "stop" | "tool_calls" | "length" | "content_filter" | "error" }
  | {
      kind: "usage";
      input?: number;
      output?: number;
      cacheRead?: number;
      cacheWrite?: number;
      modelUid?: string;
    };

function encodeImageData(img: { mimeType?: string; base64Data?: string }): Buffer {
  return Buffer.concat([
    encodeString(1, img.base64Data ?? ""),
    encodeString(2, img.mimeType ?? "image/png"),
  ]);
}

function encodeChatToolCall(tc: { id: string; name: string; arguments: string }): Buffer {
  return Buffer.concat([encodeString(1, tc.id), encodeString(2, tc.name), encodeString(3, tc.arguments)]);
}

export function encodeChatMessagePrompt(item: ChatHistoryItem, messageId: string): Buffer {
  const text = item.content
    .filter((part) => part.type === "text")
    .map((part) => part.text ?? "")
    .join("\n");
  const parts: Buffer[] = [
    encodeString(1, messageId),
    encodeVarintField(2, SOURCE_BY_ROLE[item.role] ?? 1),
    encodeString(3, text),
  ];
  for (const tc of item.tool_calls ?? []) parts.push(encodeMessage(6, encodeChatToolCall(tc)));
  if (item.tool_call_id) parts.push(encodeString(7, item.tool_call_id));
  if (item.tool_result_is_error) parts.push(encodeVarintField(9, 1));
  for (const img of item.content.filter((part) => part.type === "image")) {
    parts.push(encodeMessage(10, encodeImageData(img)));
  }
  if (item.thinking) parts.push(encodeString(11, item.thinking));
  if (item.signature) parts.push(encodeString(12, item.signature));
  if (item.thinking_redacted) parts.push(encodeVarintField(13, 1));
  return Buffer.concat(parts);
}

function encodeCompletionConfiguration(args: {
  maxTokens: number;
  temperature: number;
  topP: number;
  topK: number;
}): Buffer {
  return Buffer.concat([
    encodeVarintField(1, 1),
    encodeVarintField(2, args.maxTokens),
    encodeVarintField(3, 400),
    encodeFixed64Field(5, args.temperature),
    encodeVarintField(7, args.topK),
    encodeFixed64Field(8, args.topP),
  ]);
}

function encodeToolDef(tool: ToolDef): Buffer {
  const description = tool.description.length > 6_998 ? `${tool.description.slice(0, 6_995)}...` : tool.description;
  return Buffer.concat([
    encodeString(1, tool.name),
    encodeString(2, description),
    encodeString(3, JSON.stringify(tool.parameters ?? {})),
  ]);
}

export interface WireRequestArgs {
  apiKey: string;
  userJwt: string;
  modelUid: string;
  systemPrompt?: string;
  messages: ChatHistoryItem[];
  tools?: ToolDef[];
  cascadeId: string;
  promptId: string;
  sessionId: string;
  requestId: bigint;
  triggerId: string;
  maxTokens: number;
  temperature: number;
  topP: number;
  topK: number;
}

export function buildGetChatMessageRequest(args: WireRequestArgs): { body: Buffer; historyBytes: number } {
  const metadata = buildMetadata({
    apiKey: args.apiKey,
    userJwt: args.userJwt,
    sessionId: args.sessionId,
    requestId: args.requestId,
    triggerId: args.triggerId,
  });
  const promptFields: Buffer[] = [];
  args.messages.forEach((message, index) => {
    const messageId = message.message_id ?? `${args.cascadeId}-${index}`;
    promptFields.push(encodeMessage(3, encodeChatMessagePrompt(message, messageId)));
  });
  const history = Buffer.concat(promptFields);
  const body = Buffer.concat([
    encodeMessage(1, metadata),
    args.systemPrompt ? encodeString(2, args.systemPrompt) : Buffer.alloc(0),
    history,
    encodeVarintField(7, 5),
    encodeMessage(8, encodeCompletionConfiguration(args)),
    ...(args.tools ?? []).map((tool) => encodeMessage(10, encodeToolDef(tool))),
    encodeString(16, args.cascadeId),
    encodeString(17, args.promptId),
    encodeVarintField(20, 1),
    encodeString(21, args.modelUid),
  ]);
  return { body, historyBytes: history.length };
}

export function* decodeChatFrame(proto: Buffer): Generator<CloudChatEvent> {
  for (const field of iterFields(proto)) {
    if (field.num === 1 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const id = field.value.toString("utf8");
      if (id) yield { kind: "message_id", id };
    } else if (field.num === 3 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const text = field.value.toString("utf8");
      if (text) yield { kind: "text", text };
    } else if (field.num === 9 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const text = field.value.toString("utf8");
      if (text) yield { kind: "reasoning", text };
    } else if (field.num === 10 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const signature = field.value.toString("utf8");
      if (signature) yield { kind: "signature", text: signature };
    } else if (field.num === 11 && field.wire === 0 && field.value) {
      yield { kind: "thinking_redacted" };
    } else if (field.num === 23 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const uid = field.value.toString("utf8");
      if (uid) yield { kind: "actual_model", uid };
    } else if (field.num === 6 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      let id: string | undefined;
      let name: string | undefined;
      let argsDelta: string | undefined;
      for (const inner of iterFields(field.value)) {
        if (inner.wire === 2 && Buffer.isBuffer(inner.value)) {
          const text = inner.value.toString("utf8");
          if (inner.num === 1) id = text;
          else if (inner.num === 2) name = text;
          else if (inner.num === 3) argsDelta = text;
        }
      }
      if (id !== undefined && name !== undefined) yield { kind: "tool_call_start", id, name };
      if (argsDelta !== undefined) yield { kind: "tool_call_args", argsDelta, ...(id ? { id } : {}) };
    } else if (field.num === 5 && field.wire === 0) {
      const value = Number(field.value);
      let reason: Extract<CloudChatEvent, { kind: "finish" }>["reason"] = "stop";
      if (value === 10) reason = "tool_calls";
      else if (value === 1 || value === 3 || value === 5) reason = "length";
      else if (value === 11) reason = "content_filter";
      else if (value === 12 || value === 13) reason = "error";
      yield { kind: "finish", reason };
    } else if (field.num === 7 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      const usage = decodeUsageStats(field.value);
      if (usage) yield usage;
    }
  }
}

function decodeUsageStats(buf: Buffer): CloudChatEvent | null {
  let input: number | undefined;
  let output: number | undefined;
  let cacheWrite: number | undefined;
  let cacheRead: number | undefined;
  let modelUid: string | undefined;
  for (const field of iterFields(buf)) {
    if (field.wire === 0) {
      const n = Number(field.value);
      if (field.num === 2) input = n;
      else if (field.num === 3) output = n;
      else if (field.num === 4) cacheWrite = n;
      else if (field.num === 5) cacheRead = n;
    } else if (field.num === 9 && field.wire === 2 && Buffer.isBuffer(field.value)) {
      modelUid = field.value.toString("utf8");
    }
  }
  if (input === undefined && output === undefined && cacheRead === undefined && cacheWrite === undefined) {
    return null;
  }
  return { kind: "usage", input, output, cacheRead, cacheWrite, modelUid };
}

export interface TrailerError {
  code: string;
  message: string;
}

export function readTrailerError(text: string): TrailerError | null {
  if (!text.includes('"error"')) return null;
  try {
    const parsed = JSON.parse(text) as { error?: { code?: unknown; message?: unknown } };
    if (!parsed.error || typeof parsed.error !== "object") return null;
    return {
      code: typeof parsed.error.code === "string" ? parsed.error.code : "",
      message: typeof parsed.error.message === "string" ? parsed.error.message : "",
    };
  } catch {
    return { code: "", message: text.slice(0, 512) };
  }
}

const sessionCache = new Map<string, { sessionId: string; cascadeId: string }>();

export function sessionIds(
  apiKey: string,
  host: string,
  options?: { sessionId?: string; cacheRetention?: string },
): { sessionId: string; cascadeId: string } {
  // Compaction and other no-cache callers must not share or reuse a cascade.
  if (options?.cacheRetention === "none") {
    return { sessionId: randomUUID(), cascadeId: randomUUID() };
  }
  // Pi threads its session id through so each session tree keeps its own
  // Devin cascade instead of sharing one process-wide cascade.
  if (options?.sessionId) {
    return { sessionId: options.sessionId, cascadeId: options.sessionId };
  }
  const key = `${host}\x1f${apiKey}`;
  let ids = sessionCache.get(key);
  if (!ids) {
    ids = { sessionId: randomUUID(), cascadeId: randomUUID() };
    sessionCache.set(key, ids);
  }
  return ids;
}

function mergeHeaders(extra: Record<string, string | null> | undefined): Record<string, string> {
  const headers = new Map<string, string>([
    ["content-type", "application/connect+proto"],
    ["connect-protocol-version", "1"],
    ["connect-accept-encoding", "gzip"],
    ["user-agent", "connect-go/1.18.1 (go1.26.3)"],
  ]);
  for (const [name, value] of Object.entries(extra ?? {})) {
    if (value == null) headers.delete(name.toLowerCase());
    else headers.set(name.toLowerCase(), value);
  }
  return Object.fromEntries(headers);
}

async function* streamChatEvents(args: {
  host: string;
  request: { body: Buffer; historyBytes: number };
  headers?: Record<string, string | null>;
  signal?: AbortSignal;
  fetchFn?: typeof fetch;
  onResponse?: (response: { status: number; headers: Record<string, string> }) => void | Promise<void>;
}): AsyncGenerator<CloudChatEvent> {
  const host = args.host.replace(/\/$/, "");
  const resp = await (args.fetchFn ?? fetch)(
    `${host}/exa.api_server_pb.ApiServerService/GetChatMessage`,
    {
      method: "POST",
      headers: mergeHeaders(args.headers),
      body: new Uint8Array(frameConnectStream(args.request.body, false)),
      signal: args.signal,
    },
  );
  if (args.onResponse) {
    await args.onResponse({ status: resp.status, headers: Object.fromEntries(resp.headers.entries()) });
  }
  if (!resp.ok) {
    throw new Error(`GetChatMessage HTTP ${resp.status}: ${(await resp.text()).slice(0, 300)}`);
  }
  if (!resp.body) throw new Error("GetChatMessage returned an empty body");

  const reader = resp.body.getReader();
  const queue: Buffer[] = [];
  let queued = 0;
  let sawEos = false;
  let producedOutput = false;
  let trailerError: TrailerError | null = null;

  const peek = (n: number): Buffer | null => {
    if (queued < n) return null;
    if (queue.length === 1 && queue[0].length >= n) return queue[0].subarray(0, n);
    const parts: Buffer[] = [];
    let remaining = n;
    for (const chunk of queue) {
      if (remaining <= 0) break;
      if (chunk.length <= remaining) {
        parts.push(chunk);
        remaining -= chunk.length;
      } else {
        parts.push(chunk.subarray(0, remaining));
        remaining = 0;
      }
    }
    return Buffer.concat(parts, n);
  };

  const drop = (n: number): void => {
    queued -= n;
    let remaining = n;
    while (remaining > 0 && queue.length > 0) {
      const head = queue[0];
      if (head.length <= remaining) {
        queue.shift();
        remaining -= head.length;
      } else {
        queue[0] = head.subarray(remaining);
        remaining = 0;
      }
    }
  };

  try {
    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      if (value) {
        queue.push(Buffer.from(value));
        queued += value.length;
      }
      while (queued >= 5) {
        const header = peek(5);
        if (!header) break;
        const flags = header[0];
        const len = header.readUInt32BE(1);
        if (len > MAX_CONNECT_FRAME_PAYLOAD) {
          throw new Error(`Devin stream frame too large: ${len} bytes`);
        }
        if (queued < 5 + len) break;
        drop(5);
        const raw = peek(len) ?? Buffer.alloc(0);
        drop(len);
        let payload = raw;
        if (flags & 0x01) payload = zlib.gunzipSync(raw);
        if (flags & 0x02) {
          sawEos = true;
          trailerError = readTrailerError(payload.toString("utf8").trim());
          continue;
        }
        for (const event of decodeChatFrame(payload)) {
          if (event.kind === "text" || event.kind === "reasoning" || event.kind === "tool_call_start") {
            producedOutput = true;
          }
          yield event;
        }
      }
    }
  } finally {
    try {
      reader.releaseLock();
    } catch {
      // ignore
    }
    try {
      void resp.body?.cancel();
    } catch {
      // ignore
    }
  }

  if (trailerError) {
    if (
      !producedOutput &&
      trailerError.code.toLowerCase() === "invalid_argument" &&
      /\binternal error\b/i.test(trailerError.message) &&
      args.request.historyBytes >= LARGE_HISTORY_RECOVERY_BYTES
    ) {
      throw new Error(
        `Devin stream rejected: ${trailerError.message} (message history of ${args.request.historyBytes} bytes likely exceeds the context window)`,
      );
    }
    throw new Error(
      trailerError.code
        ? `Devin stream rejected (${trailerError.code}): ${trailerError.message}`
        : `Devin stream rejected: ${trailerError.message}`,
    );
  }
  if (!sawEos) throw new Error("Devin stream ended without an EOS trailer");
}

export function streamDevin(
  model: Model<Api>,
  context: Context,
  options?: SimpleStreamOptions,
): AssistantMessageEventStream {
  const stream = createAssistantMessageEventStream();

  void (async () => {
    const output: AssistantMessage = {
      role: "assistant",
      content: [],
      api: model.api,
      provider: model.provider,
      model: model.id,
      usage: {
        input: 0,
        output: 0,
        cacheRead: 0,
        cacheWrite: 0,
        totalTokens: 0,
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      },
      stopReason: "stop",
      timestamp: Date.now(),
    };

    let textOpen = false;
    let thinkingOpen = false;
    const openTools = new Map<string, number>();
    const toolJson = new Map<string, string>();
    let lastToolId: string | undefined;

    const closeText = () => {
      if (!textOpen) return;
      const idx = output.content.length - 1;
      const block = output.content[idx];
      if (block.type === "text") {
        stream.push({ type: "text_end", contentIndex: idx, content: block.text, partial: output });
      }
      textOpen = false;
    };
    const closeThinking = () => {
      if (!thinkingOpen) return;
      const idx = output.content.length - 1;
      const block = output.content[idx];
      if (block.type === "thinking") {
        stream.push({ type: "thinking_end", contentIndex: idx, content: block.thinking, partial: output });
      }
      thinkingOpen = false;
    };
    const closeTool = (id: string) => {
      const idx = openTools.get(id);
      if (idx === undefined) return;
      openTools.delete(id);
      const block = output.content[idx];
      if (block.type === "toolCall") {
        const json = toolJson.get(id) ?? "";
        try {
          block.arguments = JSON.parse(json);
        } catch {
          // keep last parsed object
        }
        stream.push({
          type: "toolcall_end",
          contentIndex: idx,
          toolCall: { type: "toolCall", id, name: block.name, arguments: block.arguments },
          partial: output,
        });
      }
      toolJson.delete(id);
      if (lastToolId === id) lastToolId = undefined;
    };
    const closeAllTools = () => {
      for (const id of [...openTools.keys()]) closeTool(id);
    };
    const currentThinkingBlock = () => {
      const block = output.content[output.content.length - 1];
      return block?.type === "thinking" ? block : undefined;
    };

    try {
      const apiKey = options?.apiKey;
      if (!apiKey) throw new Error("No Devin credentials. Run /login devin (uses the local Devin CLI).");
      if (options?.deferred) throw new Error("The devin provider does not support deferred responses");
      const host = (options?.env?.DEVIN_API_SERVER_URL || "https://server.codeium.com").replace(/\/$/, "");
      const modelUid = resolveModelUid(model.id, model.thinkingLevelMap, options?.reasoning);
      output.providerThinkingLevel = modelUid;
      output.responseModel = modelUid;
      const mapped = mapContextToChat(context, model);
      const timeoutMs = options?.timeoutMs;
      const signal =
        timeoutMs && timeoutMs > 0
          ? anySignal(
              [options?.signal, AbortSignal.timeout(timeoutMs)].filter((s): s is AbortSignal => Boolean(s)),
            )
          : options?.signal;
      const ids = sessionIds(apiKey, host, options);
      const userJwt = await getCachedUserJwt(apiKey, host, signal);

      const sampling =
        model.samplingParams || options?.samplingParams
          ? { ...model.samplingParams, ...options?.samplingParams }
          : {};
      const wireArgs = {
        apiKey,
        userJwt,
        modelUid,
        systemPrompt: mapped.systemPrompt,
        messages: mapped.messages,
        tools: mapped.tools.length > 0 ? mapped.tools : undefined,
        cascadeId: ids.cascadeId,
        promptId: randomUUID(),
        sessionId: ids.sessionId,
        requestId: BigInt(Date.now()),
        triggerId: randomUUID(),
        maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens),
        temperature: options?.temperature ?? 1.0,
        topP: typeof sampling.topP === "number" ? sampling.topP : 0.95,
        topK: typeof sampling.topK === "number" ? sampling.topK : 40,
      };
      const overridden = await options?.onPayload?.(wireArgs, model);
      const finalArgs =
        overridden && typeof overridden === "object" ? { ...wireArgs, ...overridden } : wireArgs;
      const request = buildGetChatMessageRequest(finalArgs);

      stream.push({ type: "start", partial: output });

      for await (const event of streamChatEvents({
        host,
        request,
        headers: options?.headers,
        signal,
        fetchFn: options?.fetch,
        onResponse: options?.onResponse
          ? (response) => options?.onResponse?.(response, model)
          : undefined,
      })) {
        if (event.kind === "text") {
          closeThinking();
          if (!textOpen) {
            output.content.push({ type: "text", text: "" });
            stream.push({ type: "text_start", contentIndex: output.content.length - 1, partial: output });
            textOpen = true;
          }
          const idx = output.content.length - 1;
          const block = output.content[idx];
          if (block.type === "text") {
            block.text += event.text;
            stream.push({ type: "text_delta", contentIndex: idx, delta: event.text, partial: output });
          }
        } else if (event.kind === "reasoning") {
          closeText();
          if (!thinkingOpen) {
            output.content.push({ type: "thinking", thinking: "" });
            stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output });
            thinkingOpen = true;
          }
          const idx = output.content.length - 1;
          const block = output.content[idx];
          if (block.type === "thinking") {
            block.thinking += event.text;
            stream.push({ type: "thinking_delta", contentIndex: idx, delta: event.text, partial: output });
          }
        } else if (event.kind === "signature") {
          const block = currentThinkingBlock();
          if (block) block.thinkingSignature = event.text;
        } else if (event.kind === "thinking_redacted") {
          const block = currentThinkingBlock();
          if (block) block.redacted = true;
        } else if (event.kind === "message_id") {
          if (!output.responseId) output.responseId = event.id;
        } else if (event.kind === "actual_model") {
          output.responseModel = event.uid;
        } else if (event.kind === "tool_call_start") {
          closeText();
          closeThinking();
          output.content.push({ type: "toolCall", id: event.id, name: event.name, arguments: {} });
          openTools.set(event.id, output.content.length - 1);
          toolJson.set(event.id, "");
          lastToolId = event.id;
          stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output });
        } else if (event.kind === "tool_call_args") {
          const id = event.id ?? lastToolId;
          const idx = id === undefined ? undefined : openTools.get(id);
          if (id === undefined || idx === undefined) continue;
          const json = (toolJson.get(id) ?? "") + event.argsDelta;
          toolJson.set(id, json);
          const block = output.content[idx];
          if (block.type === "toolCall") {
            try {
              block.arguments = JSON.parse(json);
            } catch {
              // incomplete json
            }
          }
          stream.push({ type: "toolcall_delta", contentIndex: idx, delta: event.argsDelta, partial: output });
        } else if (event.kind === "finish") {
          closeText();
          closeThinking();
          closeAllTools();
          if (event.reason === "content_filter") {
            throw new Error("Devin stop_reason: content_filter");
          }
          if (event.reason === "error") {
            throw new Error("Devin stop_reason: error");
          }
          output.stopReason =
            event.reason === "tool_calls" ? "toolUse" : event.reason === "length" ? "length" : "stop";
        } else if (event.kind === "usage") {
          output.usage.input = event.input ?? 0;
          output.usage.output = event.output ?? 0;
          output.usage.cacheRead = event.cacheRead ?? 0;
          output.usage.cacheWrite = event.cacheWrite ?? 0;
          output.usage.totalTokens =
            output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
          if (!output.responseModel && event.modelUid) output.responseModel = event.modelUid;
          calculateCost(model, output.usage);
        }
      }

      closeText();
      closeThinking();
      closeAllTools();
      stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output });
      stream.end();
    } catch (error) {
      output.stopReason = options?.signal?.aborted ? "aborted" : "error";
      output.errorMessage = error instanceof Error ? error.message : String(error);
      stream.push({ type: "error", reason: output.stopReason as "aborted" | "error", error: output });
      stream.end();
    }
  })();

  return stream;
}