import { Worker } from "node:worker_threads";
import { writeChunks } from "./chunkStorage";
import { formatForPath } from "./formats";
import { appendLogEntry, logger } from "./logger";
import type { CapturedLog } from "./schemas";
import { getSessionProcess } from "./sessionProcess";
import { finalizeLogUpdate } from "./store";

type BaseFinalizeLogJob = {
  log: CapturedLog;
  upstreamUrl: string;
  elapsedMs: number;
};

export type FinalizeNonStreamingLogJob = BaseFinalizeLogJob & {
  type: "non-streaming";
  responseStatus: number;
  responseBody: string;
};

export type FinalizeStreamingLogJob = BaseFinalizeLogJob & {
  type: "streaming";
  responseStatus: number;
  rawStream: string;
};

export type FinalizeStreamAbortLogJob = BaseFinalizeLogJob & {
  type: "stream-abort";
  rawStream: string;
  hasChunks: boolean;
};

export type FinalizeLogJob =
  | FinalizeNonStreamingLogJob
  | FinalizeStreamingLogJob
  | FinalizeStreamAbortLogJob;

export type FinalizeLogResult = {
  log: CapturedLog;
  upstreamUrl: string;
  error: string | null;
};

export function buildFileLogEntry(log: CapturedLog, upstreamUrl: string): Record<string, unknown> {
  return {
    timestamp: log.timestamp,
    id: log.id,
    method: log.method,
    path: log.path,
    model: log.model,
    sessionId: log.sessionId,
    rawRequestBody: log.rawRequestBody,
    responseStatus: log.responseStatus,
    responseText: log.responseText,
    inputTokens: log.inputTokens,
    outputTokens: log.outputTokens,
    cacheCreationInputTokens: log.cacheCreationInputTokens,
    cacheReadInputTokens: log.cacheReadInputTokens,
    elapsedMs: log.elapsedMs,
    streaming: log.streaming,
    userAgent: log.userAgent,
    origin: log.origin,
    upstreamUrl,
    clientPort: log.clientPort,
    clientPid: log.clientPid,
    clientCwd: log.clientCwd,
    clientProjectFolder: log.clientProjectFolder,
    streamingChunks: log.streamingChunks,
    streamingChunksPath: log.streamingChunksPath,
    error: log.error,
  };
}

function cloneLog(log: CapturedLog): CapturedLog {
  return { ...log };
}

function errorMessage(err: unknown): string {
  return err instanceof Error ? err.message : String(err);
}

function persistStreamingChunks(log: CapturedLog): void {
  if (log.streamingChunks !== undefined && log.streamingChunks.chunks.length > 0) {
    const chunkPath = writeChunks(
      log.id,
      log.streamingChunks.chunks,
      log.streamingChunks.truncated,
    );
    log.streamingChunksPath = chunkPath;
  }
}

function finalizeWithError(
  job: FinalizeLogJob,
  log: CapturedLog,
  fallbackStatus: number,
  fallbackResponseText: string,
  err: unknown,
): FinalizeLogResult {
  const message = errorMessage(err);
  logger.error(`[logFinalizer] Failed to finalize log #${log.id}:`, message);
  log.responseStatus = log.responseStatus ?? fallbackStatus;
  log.responseText = log.responseText ?? fallbackResponseText;
  log.elapsedMs = job.elapsedMs;
  log.error = message;
  return { log, upstreamUrl: job.upstreamUrl, error: message };
}

function finalizeNonStreaming(
  job: FinalizeNonStreamingLogJob,
  log: CapturedLog,
): FinalizeLogResult {
  const formatHandler = formatForPath(log.path);
  if (formatHandler === null) {
    return finalizeWithError(job, log, job.responseStatus, job.responseBody, "Unsupported format");
  }

  try {
    const tokens = formatHandler.extractTokens(job.responseBody);

    log.elapsedMs = job.elapsedMs;
    log.responseStatus = job.responseStatus;
    log.responseText = job.responseBody;
    log.inputTokens = tokens.inputTokens;
    log.outputTokens = tokens.outputTokens;
    log.cacheCreationInputTokens = tokens.cacheCreationInputTokens;
    log.cacheReadInputTokens = tokens.cacheReadInputTokens;

    return { log, upstreamUrl: job.upstreamUrl, error: null };
  } catch (err) {
    return finalizeWithError(job, log, job.responseStatus, job.responseBody, err);
  }
}

function finalizeStreaming(job: FinalizeStreamingLogJob, log: CapturedLog): FinalizeLogResult {
  const formatHandler = formatForPath(log.path);
  if (formatHandler === null) {
    return finalizeWithError(job, log, job.responseStatus, job.rawStream, "Unsupported format");
  }

  try {
    log.elapsedMs = job.elapsedMs;
    log.responseStatus = job.responseStatus;
    log.responseText = formatHandler.extractStream(
      job.rawStream,
      log,
      log.model ?? undefined,
      true,
    );
    persistStreamingChunks(log);

    return { log, upstreamUrl: job.upstreamUrl, error: null };
  } catch (err) {
    return finalizeWithError(job, log, job.responseStatus, job.rawStream, err);
  }
}

function finalizeStreamAbort(job: FinalizeStreamAbortLogJob, log: CapturedLog): FinalizeLogResult {
  const formatHandler = formatForPath(log.path);
  if (formatHandler === null && job.hasChunks) {
    return finalizeWithError(job, log, 499, "Client aborted", "Unsupported format");
  }

  try {
    log.elapsedMs = job.elapsedMs;
    if (job.hasChunks && formatHandler !== null) {
      log.responseText = formatHandler.extractStream(
        job.rawStream,
        log,
        log.model ?? undefined,
        true,
      );
      persistStreamingChunks(log);
    } else {
      log.responseText = "Client aborted";
    }

    return { log, upstreamUrl: job.upstreamUrl, error: "Client aborted" };
  } catch (err) {
    return finalizeWithError(job, log, 499, "Client aborted", err);
  }
}

export function buildFinalizeLogResult(job: FinalizeLogJob): FinalizeLogResult {
  const log = cloneLog(job.log);

  switch (job.type) {
    case "non-streaming":
      return finalizeNonStreaming(job, log);
    case "streaming":
      return finalizeStreaming(job, log);
    case "stream-abort":
      return finalizeStreamAbort(job, log);
  }
}

export function commitFinalizeLogResult(result: FinalizeLogResult): void {
  appendLogEntry({ ...buildFileLogEntry(result.log, result.upstreamUrl), error: result.error });
  finalizeLogUpdate(result.log);
}

// ── Worker pool ────────────────────────────────────────────────
// buildFinalizeLogResult is the CPU-heavy step (token extraction,
// SSE parsing, chunk persistence).  We offload it to a fixed pool
// of Worker Threads so the main thread stays responsive even under
// concurrent sessions.  commitFinalizeLogResult stays in-process
// because it touches the in-memory cache, SSE emitter, and
// append‑only log — all of which belong to the main thread.

// ── Routing ─────────────────────────────────────────────────────
// FINALIZER_RUNTIME selects the execution backend:
//   "process" → per-session child process (default, max isolation)
//   "worker"  → shared Worker Thread pool
//   "inline"  → synchronous in-process (debug / fallback)
// For backward compatibility, FINALIZER_USE_WORKER=0 forces "inline".

const RUNTIME: "process" | "worker" | "inline" = (() => {
  if (process.env["FINALIZER_USE_WORKER"] === "0") return "inline";
  const mode = process.env["FINALIZER_RUNTIME"];
  if (mode === "worker") return "worker";
  if (mode === "inline") return "inline";
  return "process";
})();

function executeBuildInSessionProcess(job: FinalizeLogJob): Promise<FinalizeLogResult> {
  const sessionId = job.log.sessionId ?? "__unassigned__";
  return getSessionProcess(sessionId)
    .enqueue(job)
    .catch((err: unknown) => {
      const message = err instanceof Error ? err.message : String(err);
      logger.error(
        `[logFinalizer] Session process failed for log #${job.log.id}, ` +
          `falling back to in-process: ${message}`,
      );
      return buildFinalizeLogResult(job);
    });
}

function resolveBuildPromise(job: FinalizeLogJob): Promise<FinalizeLogResult> {
  switch (RUNTIME) {
    case "process":
      return executeBuildInSessionProcess(job);
    case "worker":
      return executeBuildInWorker(job);
    case "inline":
      return Promise.resolve(buildFinalizeLogResult(job));
  }
}

const WORKER_COUNT = Math.max(1, Number(process.env["FINALIZER_WORKER_COUNT"]) || 4);
let _workers: Worker[] | null = null;
const _pending = new Map<string, (result: FinalizeLogResult) => void>();
let _nextWorker = 0;
let _jobSeq = 0;

function getWorkers(): Worker[] {
  if (_workers !== null) return _workers;
  _workers = [];
  for (let i = 0; i < WORKER_COUNT; i++) {
    const w = new Worker(new URL("./logFinalizer.worker.ts", import.meta.url));
    w.on("message", (msg: { id: string; result: FinalizeLogResult }) => {
      const resolve = _pending.get(msg.id);
      if (resolve !== undefined) {
        _pending.delete(msg.id);
        resolve(msg.result);
      }
    });
    w.on("error", (err) => {
      logger.error(
        "[logFinalizer] Worker error:",
        err instanceof Error ? err.message : String(err),
      );
    });
    _workers.push(w);
  }
  return _workers;
}

export function executeBuildInWorker(job: FinalizeLogJob): Promise<FinalizeLogResult> {
  return new Promise((resolve) => {
    const list = getWorkers();
    // getWorkers() always returns at least 1 worker; the array index is safe
    const idx = _nextWorker % list.length;
    _nextWorker++;
    const w = list[idx];
    if (w === undefined) return; // unreachable, satisfies type-checker
    const id = String(++_jobSeq);
    _pending.set(id, resolve);
    w.postMessage({ id, job });
  });
}

export function executeFinalizeLogJob(job: FinalizeLogJob): Promise<void> {
  return resolveBuildPromise(job).then((result) => {
    commitFinalizeLogResult(result);
    if (result.error !== null && result.error !== "Client aborted") {
      return Promise.reject(new Error(result.error));
    }
    return undefined;
  });
}
