// src/codec/rtkRouteStream.ts
//
// TransformStream that rewrites `command` fields in Bash tool_use SSE events
// so they are prefixed with `rtk ` when the first token is an RTK-routed command.
//
// State machine: IDLE → DETECTING → INJECTED / PASSTHROUGH_AS_IS → IDLE
//   IDLE:            pass events through; enter DETECTING on Bash tool_use with empty input
//   DETECTING:       buffer content_block_delta events; scan accumulated partial_json for
//                    "command" key then read first token; on RTK-routable → INJECTED;
//                    on non-RTK → PASSTHROUGH_AS_IS; on stop without detection → flush+IDLE
//   INJECTED:        start event already streamed live; one synthetic delta emitted; pass through
//   PASSTHROUGH_AS_IS: start event already streamed live; buffered deltas flushed; pass through

import { encodeEvent, parseEvent, splitEvents } from "./sseParser.ts";

export const RTK_COMMANDS = new Set([
  "git", "gh",
  "pnpm", "npm", "npx", "yarn",
  "find", "grep", "ls", "tree",
  "tsc", "jest", "vitest", "playwright", "lint", "prettier", "next", "prisma", "dotnet", "cargo",
  "curl", "wget",
  "docker", "kubectl", "psql", "aws",
  "diff", "wc",
]);

/**
 * Pure function. Rewrites a shell command string to go through RTK if the first
 * token is in RTK_COMMANDS. Returns the original string if no rewrite needed.
 */
export function rewriteCommand(cmd: string): string {
  const trimmed = cmd.trimStart();
  const spaceIdx = trimmed.search(/\s/);
  const firstToken = spaceIdx === -1 ? trimmed : trimmed.slice(0, spaceIdx);
  if (RTK_COMMANDS.has(firstToken)) {
    return "rtk " + trimmed;
  }
  return cmd;
}

type State = "IDLE" | "DETECTING" | "INJECTED" | "PASSTHROUGH_AS_IS";

/**
 * Scan accumulated partial_json for the "command" JSON key, then read
 * the first token of the value (up to space/quote/backslash/EOF).
 *
 * Returns:
 *   { found: true, token: string, tokenStartIdx: number } — first token found
 *   { found: false, incomplete: true } — key found but value not buffered yet (EOF before first token char)
 *   { found: false, incomplete: false } — key not found in accumulated string
 *
 * The regex is reset to lastIndex=0 on every call to ensure reentrant use.
 * Uses matchAll (g flag) to iterate over all occurrences.
 */
const COMMAND_KEY_RE = /[{,]\s*"command"\s*:\s*"/g;

type DetectResult =
  | { found: true; token: string; tokenStartIdx: number }
  | { found: false; incomplete: boolean };

function detectFirstToken(accumulated: string): DetectResult {
  COMMAND_KEY_RE.lastIndex = 0;

  for (const match of accumulated.matchAll(COMMAND_KEY_RE)) {
    // Skip escaped key: if the char before match is backslash (0x5C), this is
    // \"command\" inside a string value, not a real key. Skip it.
    if (match.index !== undefined && match.index > 0 && accumulated.charCodeAt(match.index - 1) === 0x5c) {
      continue;
    }

    // match[0] ends with `"`, so value starts right after match.index + match[0].length
    const matchStart = match.index ?? 0;
    const valueStart = matchStart + match[0].length;

    if (valueStart >= accumulated.length) {
      // Nothing after the opening quote yet
      return { found: false, incomplete: true };
    }

    // Read first token: terminates at space (0x20), quote (0x22), backslash (0x5C), or EOF
    const tokenStartIdx = valueStart;
    let tokenEndIdx = valueStart;
    while (tokenEndIdx < accumulated.length) {
      const ch = accumulated.charCodeAt(tokenEndIdx);
      if (ch === 0x20 || ch === 0x22 || ch === 0x5c) {
        break;
      }
      tokenEndIdx++;
    }

    if (tokenEndIdx === accumulated.length) {
      // Reached EOF without finding a terminator → incomplete, keep buffering
      return { found: false, incomplete: true };
    }

    const token = accumulated.slice(tokenStartIdx, tokenEndIdx);
    return { found: true, token, tokenStartIdx };
  }

  // Key not found yet
  return { found: false, incomplete: false };
}

/**
 * True iff the raw SSE event carries a `data:` line.
 */
function hasDataLine(rawEvent: string): boolean {
  return rawEvent.replace(/\r\n/g, "\n").split("\n").some((line) => line.startsWith("data:"));
}

/**
 * Returns a TransformStream<Uint8Array, Uint8Array> that intercepts Bash tool_use
 * content_block_start events and rewrites the `command` field through RTK.
 *
 * Two paths:
 *   Path A (input present in start event): rewrite in-place, emit immediately.
 *   Path B (input is empty object): stream start live, detect first token in deltas,
 *          emit synthetic delta on RTK match or flush buffered on non-match.
 */
export function rtkRouteStream(): TransformStream<Uint8Array, Uint8Array> {
  const enc = new TextEncoder();
  const dec = new TextDecoder("utf-8", { fatal: false });

  let state: State = "IDLE";
  let carry = "";

  // DETECTING state accumulation
  let bufferedDeltaEvents: string[] = []; // raw SSE delta event strings buffered during DETECTING
  let accumulatedPartialJson = "";
  let blockIndex: number | null = null; // the `index` field from content_block_start

  function resetDetecting(): void {
    state = "IDLE";
    bufferedDeltaEvents = [];
    accumulatedPartialJson = "";
    blockIndex = null;
  }

  /**
   * Emit a list of pre-built SSE event strings to the controller.
   */
  function emitAll(
    lines: string[],
    controller: TransformStreamDefaultController<Uint8Array>,
  ): void {
    for (const line of lines) {
      if (line.length > 0) controller.enqueue(enc.encode(line));
    }
  }

  /**
   * Quick helper to read the `type` field from a raw data event.
   * Returns empty string on failure.
   */
  function evTypeOf(rawEvent: string): string {
    const parsed = parseEvent(rawEvent).dataJson as Record<string, unknown> | null;
    if (parsed === null) return "";
    return typeof parsed.type === "string" ? parsed.type : "";
  }

  /**
   * Attempt detection on accumulated partial JSON and emit/transition accordingly.
   * Called after each new delta is appended to accumulatedPartialJson.
   * Side-effect: may emit a synthetic delta (INJECTED) or flush buffered deltas (PASSTHROUGH_AS_IS).
   */
  function processStreaming(
    controller: TransformStreamDefaultController<Uint8Array>,
  ): void {
    const result = detectFirstToken(accumulatedPartialJson);

    if (!result.found) {
      // Incomplete or key not found — keep buffering
      return;
    }

    const { token, tokenStartIdx } = result;

    if (RTK_COMMANDS.has(token)) {
      // RTK-routable: emit ONE synthetic delta with rewritten partial JSON, drop buffered deltas
      const idx = blockIndex !== null ? blockIndex : 0;
      const rewrittenPartial =
        accumulatedPartialJson.slice(0, tokenStartIdx) +
        "rtk " +
        accumulatedPartialJson.slice(tokenStartIdx);

      const syntheticDelta = encodeEvent("content_block_delta", {
        type: "content_block_delta",
        index: idx,
        delta: {
          type: "input_json_delta",
          partial_json: rewrittenPartial,
        },
      });

      controller.enqueue(enc.encode(syntheticDelta));
      // Drop buffered delta events (not emitted — start was already live)
      bufferedDeltaEvents = [];
      accumulatedPartialJson = "";
      state = "INJECTED";
    } else {
      // Non-RTK: flush all buffered deltas as-is
      emitAll(bufferedDeltaEvents, controller);
      bufferedDeltaEvents = [];
      accumulatedPartialJson = "";
      state = "PASSTHROUGH_AS_IS";
    }
  }

  /**
   * Process a single SSE event in IDLE state.
   * Returns an array of event strings to emit.
   * Side-effect: may change `state` to DETECTING.
   */
  function processIdle(rawEvent: string): string[] {
    if (!hasDataLine(rawEvent)) {
      return [rawEvent];
    }

    const parsed = parseEvent(rawEvent).dataJson as Record<string, unknown> | null;
    if (parsed === null) {
      return [rawEvent];
    }

    if (parsed.type !== "content_block_start") {
      return [rawEvent];
    }

    const cb = parsed.content_block as Record<string, unknown> | undefined;
    if (!cb || cb.type !== "tool_use" || cb.name !== "Bash") {
      return [rawEvent];
    }

    const input = cb.input;

    // Path A: input has a command string field → rewrite in-place, emit immediately
    if (
      typeof input === "object" &&
      input !== null &&
      "command" in input &&
      typeof (input as Record<string, unknown>).command === "string"
    ) {
      const inputRec = input as Record<string, unknown>;
      const original = inputRec.command as string;
      const rewritten = rewriteCommand(original);
      if (rewritten !== original) {
        inputRec.command = rewritten;
        // Re-encode the event with preserved `event:` line.
        return [encodeEvent("content_block_start", parsed)];
      }
      return [rawEvent];
    }

    // Path B: empty input object → stream start event LIVE, enter DETECTING
    if (typeof input === "object" && input !== null && !("command" in input)) {
      state = "DETECTING";
      bufferedDeltaEvents = [];
      accumulatedPartialJson = "";
      blockIndex = typeof parsed.index === "number" ? parsed.index : null;
      // Emit the start event live — do NOT push to bufferedDeltaEvents
      return [rawEvent];
    }

    // input is null/undefined/unexpected type → pass through
    return [rawEvent];
  }

  /**
   * Process a single SSE event while in DETECTING, INJECTED, or PASSTHROUGH_AS_IS state.
   * Returns an array of event strings to emit immediately.
   * Side-effect: may flush buffered events, emit synthetic delta, and change state.
   */
  function processNonIdle(
    rawEvent: string,
    controller: TransformStreamDefaultController<Uint8Array>,
  ): string[] {
    if (state === "DETECTING") {
      if (!hasDataLine(rawEvent)) {
        // Non-data line during DETECTING: buffer it
        bufferedDeltaEvents.push(rawEvent);
        return [];
      }

      const parsed = parseEvent(rawEvent).dataJson as Record<string, unknown> | null;
      if (parsed === null) {
        bufferedDeltaEvents.push(rawEvent);
        return [];
      }

      const evType = parsed.type as string | undefined;

      if (evType === "content_block_start") {
        // Malformed stream: new block while detecting → flush buffered, handle as IDLE
        const toFlush = [...bufferedDeltaEvents];
        resetDetecting();
        emitAll(toFlush, controller);
        return processIdle(rawEvent);
      }

      if (evType === "content_block_delta") {
        const delta = parsed.delta as Record<string, unknown> | undefined;
        if (delta && delta.type === "input_json_delta" && typeof delta.partial_json === "string") {
          accumulatedPartialJson += delta.partial_json as string;
        }
        bufferedDeltaEvents.push(rawEvent);

        // Attempt detection after accumulating more data
        processStreaming(controller);
        return [];
      }

      if (evType === "content_block_stop") {
        // Stop arrived while still detecting: flush buffered deltas as-is, emit stop, → IDLE
        const toFlush = [...bufferedDeltaEvents];
        resetDetecting();
        return [...toFlush, rawEvent];
      }

      // Other event types while detecting → buffer
      bufferedDeltaEvents.push(rawEvent);
      return [];
    }

    // INJECTED or PASSTHROUGH_AS_IS: pass events through
    if (evTypeOf(rawEvent) === "content_block_stop") {
      state = "IDLE";
      blockIndex = null;
    }
    return [rawEvent];
  }

  return new TransformStream<Uint8Array, Uint8Array>({
    transform(chunk: Uint8Array, controller: TransformStreamDefaultController<Uint8Array>) {
      try {
        const chunkStr = dec.decode(chunk, { stream: true });
        const combined = carry + chunkStr;

        // Only process up to last complete event boundary
        const lastBoundary = combined.lastIndexOf("\n\n");
        if (lastBoundary === -1) {
          carry = combined;
          return;
        }

        const cutPoint = lastBoundary + 2;
        const toProcess = combined.slice(0, cutPoint);
        carry = combined.slice(cutPoint);

        const { events } = splitEvents(toProcess);

        for (const rawEvent of events) {
          let toEmit: string[];
          if (state === "IDLE") {
            toEmit = processIdle(rawEvent);
          } else {
            toEmit = processNonIdle(rawEvent, controller);
          }
          emitAll(toEmit, controller);
        }
      } catch {
        // Error safety: emit original chunk unchanged, reset state
        resetDetecting();
        carry = "";
        controller.enqueue(chunk);
      }
    },

    flush(controller: TransformStreamDefaultController<Uint8Array>) {
      try {
        // Flush decoder
        const remaining = dec.decode(new Uint8Array(0), { stream: false });
        const finalStr = carry + remaining;
        carry = "";

        if (finalStr.length > 0) {
          const { events, remainder } = splitEvents(finalStr);

          for (const rawEvent of events) {
            let toEmit: string[];
            if (state === "IDLE") {
              toEmit = processIdle(rawEvent);
            } else {
              toEmit = processNonIdle(rawEvent, controller);
            }
            emitAll(toEmit, controller);
          }

          // Any incomplete event in remainder: emit as-is
          if (remainder.length > 0) {
            if (state === "DETECTING") {
              // Flush buffered deltas + remainder
              emitAll([...bufferedDeltaEvents, remainder], controller);
              resetDetecting();
            } else {
              controller.enqueue(enc.encode(remainder));
            }
          }
        }

        // If still detecting at stream end, flush everything
        if (state === "DETECTING") {
          emitAll(bufferedDeltaEvents, controller);
          resetDetecting();
        }
      } catch {
        // Flush error: emit any carry as-is
        if (carry.length > 0) controller.enqueue(enc.encode(carry));
        if (state === "DETECTING") {
          emitAll(bufferedDeltaEvents, controller);
          resetDetecting();
        }
        carry = "";
      }
    },
  });
}