import type { MarkdownCodeBlockPreview, Message } from "#cli/core/model";
import type { ThemeId } from "#cli/core/theme";
import { Buffer } from "node:buffer";
import { appendFileSync } from "node:fs";
import { debugLog } from "#cli/core/debug-log";
import { createEnvReader, readProcessEnv } from "#cli/core/env";
import { splitMarkdownByFencedCodeBlocks } from "#cli/core/markdown";
import { resolveShikiTheme } from "#cli/core/theme";
import { isTerminalStatusText } from "#cli/i18n";
import { computeMarkdownCodeBlockPreviewWithShiki } from "../../ui/chat/read-highlight";
import { createLruCache } from "./highlight-cache";

type AssistantMessage = Extract<Message, { role: "assistant" }>;
type StatusPart = Extract<AssistantMessage["parts"][number], { type: "status" }>;
type MarkdownHighlightState = Readonly<{
  messages: Message[];
  runningAgentRound?: boolean;
  runningToolCallId?: string | null;
  ui?: {
    theme?: {
      mode?: ThemeId;
    };
  };
}>;
type HighlightJob = {
  state: MarkdownHighlightState;
  assistantId: string;
  updateAssistant: (assistantId: string, fn: (m: AssistantMessage) => AssistantMessage) => void;
};

function getTextHash(text: string): string {
  let hash = 0;
  const str = String(text ?? "");
  for (let i = 0; i < str.length; i++) {
    const ch = str.charCodeAt(i);
    hash = (hash << 5) - hash + ch;
    hash = hash & hash;
  }
  return hash.toString(36);
}

const env = createEnvReader(readProcessEnv());
const MAX_BYTES_PER_BLOCK = 200_000;
const MAX_LINES_PER_BLOCK = 2000;
const DEBOUNCE_MS = 120;
const MAX_BLOCKS_PER_MESSAGE = (() => {
  const raw = env.prefixedStr("MARKDOWN_CODE_HIGHLIGHT_MAX_BLOCKS").toLowerCase();
  // Default: highlight all fenced code blocks (no per-message budget cap).
  if (!raw || raw === "all" || raw === "0") return Number.POSITIVE_INFINITY;
  const n = Number(raw);
  if (Number.isFinite(n) && n > 0) return Math.max(1, Math.floor(n));
  return Number.POSITIVE_INFINITY;
})();

const PROFILE_MD_HL = env.prefixedStr("PROFILE_MARKDOWN_HIGHLIGHT") === "1";
const PROFILE_MD_HL_SLOW_MS = (() => {
  const raw = env.prefixedStr("PROFILE_MARKDOWN_HIGHLIGHT_SLOW_MS");
  const n = Number(raw);
  if (Number.isFinite(n) && n > 0) return Math.min(60_000, Math.max(10, Math.floor(n)));
  return 120;
})();

const blockCache = createLruCache<string, MarkdownCodeBlockPreview>(
  (() => {
    const raw = env.prefixedStr("MARKDOWN_CODE_HIGHLIGHT_CACHE");
    const n = Number(raw);
    if (Number.isFinite(n) && n > 0) return Math.min(20_000, Math.max(100, Math.floor(n)));
    return 1000;
  })(),
);

function logPerf(line: string): void {
  if (!PROFILE_MD_HL) return;
  try {
    const now = Date.now();
    const d = new Date(now);
    const hh = String(d.getHours()).padStart(2, "0");
    const mm = String(d.getMinutes()).padStart(2, "0");
    const ss = String(d.getSeconds()).padStart(2, "0");
    const mmm = String(d.getMilliseconds()).padStart(3, "0");
    const ts = `${hh}:${mm}:${ss}.${mmm}`;
    appendFileSync("/tmp/dimcode-perf.log", `[${ts}] [MD_HL] ${line}\n`);
  } catch {
    // ignore
  }
}

function isTerminalAssistantStatus(
  msg: AssistantMessage,
  state: Readonly<{
    runningAgentRound?: boolean;
    runningToolCallId?: string | null;
  }>,
): boolean {
  const status = msg.parts.find((part): part is StatusPart => part.type === "status");
  if (!status) return true; // no status => treat as stable (e.g. restored sessions)
  const text = String(status.text ?? "").trim();
  // Empty status means:
  // - restored snapshot (stable) OR
  // - currently running but has not received a status update yet (streaming)
  // Use store state as the tie-breaker.
  if (!text) return !(state.runningAgentRound || state.runningToolCallId);
  return isTerminalStatusText(text);
}

const inFlight = new Set<string>();
const pendingByAssistantId = new Map<string, ReturnType<typeof setTimeout>>();
const queuedByAssistantId = new Map<string, HighlightJob>();
const queueOrder: string[] = [];
let queueRunning = 0;
const MAX_CONCURRENCY = 1;

export function clearMarkdownCodeHighlightState(): void {
  inFlight.clear();
  pendingByAssistantId.forEach((timer) => clearTimeout(timer));
  pendingByAssistantId.clear();
  queuedByAssistantId.clear();
  queueOrder.length = 0;
  blockCache.clear();
}

function enqueueHighlightJob(job: HighlightJob): void {
  const id = String(job.assistantId ?? "").trim();
  if (!id) return;
  const existed = queuedByAssistantId.has(id);
  queuedByAssistantId.set(id, job);
  if (!existed) queueOrder.push(id);
  void drainQueue();
}

async function drainQueue(): Promise<void> {
  while (queueRunning < MAX_CONCURRENCY && queueOrder.length > 0) {
    const assistantId = queueOrder.shift()!;
    const job = queuedByAssistantId.get(assistantId);
    queuedByAssistantId.delete(assistantId);
    if (!job) continue;
    queueRunning++;
    void runHighlight(job)
      .catch(() => {})
      .finally(() => {
        queueRunning--;
        void drainQueue();
      });
  }
}

export function scheduleAssistantMarkdownCodeHighlight(
  opts: Readonly<{
    state: MarkdownHighlightState;
    assistantId: string;
    updateAssistant: (assistantId: string, fn: (m: AssistantMessage) => AssistantMessage) => void;
  }>,
): void {
  const assistantId = String(opts.assistantId ?? "").trim();
  if (!assistantId) return;

  const prevTimer = pendingByAssistantId.get(assistantId);
  if (prevTimer) clearTimeout(prevTimer);

  const timer = setTimeout(() => {
    pendingByAssistantId.delete(assistantId);
    enqueueHighlightJob({ ...opts, assistantId });
  }, DEBOUNCE_MS);
  pendingByAssistantId.set(assistantId, timer);
}

async function runHighlight(opts: Readonly<HighlightJob>): Promise<void> {
  const t0 = PROFILE_MD_HL ? Date.now() : 0;
  // Yield to let the UI paint before heavy markdown scanning/highlighting.
  await new Promise<void>((resolve) => {
    const si = (
      globalThis as typeof globalThis & {
        setImmediate?: (fn: () => void) => void;
      }
    ).setImmediate;
    if (typeof si === "function") si(() => resolve());
    else setTimeout(resolve, 0);
  });
  const msg = opts.state.messages.find((m) => m.id === opts.assistantId);
  if (!msg || msg.role !== "assistant") return;

  // During streaming, defer fenced code highlighting: scanning + hashing very large markdown
  // is a frequent source of multi-second event-loop stalls. We'll highlight once the
  // assistant status becomes terminal, or on restored sessions with no status part.
  if (
    !isTerminalAssistantStatus(msg, {
      runningAgentRound: opts.state.runningAgentRound,
      runningToolCallId: opts.state.runningToolCallId,
    })
  )
    return;

  const parts = msg.parts;
  const targets: Array<{ index: number; markdown: string }> = [];
  for (let i = 0; i < parts.length; i++) {
    const p = parts[i];
    if (p.type !== "markdown") continue;
    const md = String(p.markdown ?? "").replace(/\r/g, "");
    if (!md.includes("```")) continue;
    targets.push({ index: i, markdown: md });
  }
  if (targets.length === 0) return;

  if (targets.length > 0 && PROFILE_MD_HL)
    logPerf(`start assistant=${opts.assistantId} parts=${parts.length} targets=${targets.length}`);

  const theme = resolveShikiTheme(opts.state.ui?.theme?.mode);

  const combinedSourceHash = getTextHash(targets.map((t) => t.markdown).join("\n\u0000\n"));
  const flightKey = `${opts.assistantId}:${combinedSourceHash}:${targets.length}:${theme}`;
  if (inFlight.has(flightKey)) return;

  inFlight.add(flightKey);
  try {
    const updates = new Map<
      number,
      { sourceHash: string; codeBlocks: MarkdownCodeBlockPreview[] }
    >();
    let tokenizedBlocks = 0;

    for (const t of targets) {
      if (tokenizedBlocks >= MAX_BLOCKS_PER_MESSAGE) break;
      const blocks = splitMarkdownByFencedCodeBlocks(t.markdown);
      const codeBlocks = blocks.filter((b) => b.type === "code") as Array<
        Extract<(typeof blocks)[number], { type: "code" }>
      >;
      if (codeBlocks.length === 0) continue;

      const sourceHash = getTextHash(t.markdown);
      const part = msg.parts[t.index];
      const existing: readonly MarkdownCodeBlockPreview[] =
        part?.type === "markdown" && Array.isArray(part.codeBlocks) ? part.codeBlocks : [];

      const next: MarkdownCodeBlockPreview[] = [];
      for (let i = 0; i < codeBlocks.length; i++) {
        if (tokenizedBlocks >= MAX_BLOCKS_PER_MESSAGE) break;
        const cb = codeBlocks[i]!;
        const language = String(cb.language ?? "").trim() || "text";
        const code = String(cb.code ?? "").replace(/\r/g, "");
        const totalLines = code.split("\n").length;
        const hash = getTextHash(`${language}\0${code}`);
        const cacheKey = `${hash}:${totalLines}:${theme}`;

        const prev = existing[i];
        if (prev && prev.hash === hash && prev.totalLines === totalLines) {
          if (prev.kind === "skipped") {
            next.push(prev);
            continue;
          }
          if (prev.kind === "shiki" && String(prev.theme ?? "").trim() === theme) {
            next.push(prev);
            continue;
          }
        }

        const cached = blockCache.get(cacheKey);
        if (cached) {
          next.push(cached);
          tokenizedBlocks++;
          continue;
        }

        const bytes = Buffer.byteLength(code, "utf8");
        if (bytes > MAX_BYTES_PER_BLOCK || totalLines > MAX_LINES_PER_BLOCK) {
          const skipped = {
            kind: "skipped",
            hash,
            totalLines,
            reason: "too_large",
          } as MarkdownCodeBlockPreview;
          blockCache.set(cacheKey, skipped);
          next.push(skipped);
          continue;
        }

        const preview = await computeMarkdownCodeBlockPreviewWithShiki(code, {
          language,
          theme,
        });
        // Ensure our hash matches the scheduler's language-aware hash so stale checks stay consistent.
        if (preview.hash !== hash) {
          const normalized: MarkdownCodeBlockPreview = {
            ...preview,
            hash,
            totalLines,
          };
          blockCache.set(cacheKey, normalized);
          next.push(normalized);
        } else {
          blockCache.set(cacheKey, preview);
          next.push(preview);
        }
        tokenizedBlocks++;
      }

      updates.set(t.index, { sourceHash, codeBlocks: next });
    }

    if (updates.size === 0) return;

    opts.updateAssistant(opts.assistantId, (m) => {
      const parts = m.parts.slice();
      let changed = false;
      for (const [idx, upd] of updates.entries()) {
        const p = parts[idx];
        if (!p || p.type !== "markdown") continue;
        const nowMarkdown = String(p.markdown ?? "").replace(/\r/g, "");
        if (getTextHash(nowMarkdown) !== upd.sourceHash) continue;
        const prevBlocks = Array.isArray(p.codeBlocks) ? p.codeBlocks : null;
        const same =
          prevBlocks &&
          prevBlocks.length === upd.codeBlocks.length &&
          prevBlocks.every((block, i2) => sameMarkdownCodeBlockPreview(block, upd.codeBlocks[i2]));
        if (same) continue;
        parts[idx] = { ...p, codeBlocks: upd.codeBlocks } as const;
        changed = true;
      }
      return changed ? { ...m, parts } : m;
    });
  } catch (e) {
    debugLog("markdown_code_highlight_error", {
      assistantId: opts.assistantId,
      error: e instanceof Error ? { message: e.message, stack: e.stack } : String(e),
    });
  } finally {
    inFlight.delete(flightKey);
    if (PROFILE_MD_HL) {
      const elapsed = Date.now() - t0;
      logPerf(
        `end assistant=${opts.assistantId} ms=${elapsed} maxBlocks=${MAX_BLOCKS_PER_MESSAGE}`,
      );
      if (elapsed >= PROFILE_MD_HL_SLOW_MS)
        logPerf(
          `slow assistant=${opts.assistantId} ms=${elapsed} maxBlocks=${MAX_BLOCKS_PER_MESSAGE}`,
        );
    }
  }
}

function sameMarkdownCodeBlockPreview(
  left: MarkdownCodeBlockPreview | undefined,
  right: MarkdownCodeBlockPreview | undefined,
): boolean {
  if (!left || !right) return false;
  if (
    left.hash !== right.hash ||
    left.totalLines !== right.totalLines ||
    left.kind !== right.kind
  ) {
    return false;
  }
  if (left.kind === "shiki" && right.kind === "shiki")
    return String(left.theme ?? "").trim() === String(right.theme ?? "").trim();
  if (left.kind === "skipped" && right.kind === "skipped") return left.reason === right.reason;
  return false;
}
