/**
 * dsh-session-cost — host half.
 *
 * The harness throws away the cost its provider reports: `dsh-llm-pi-ai`'s
 * `mapUsage` copies only token counts, and `pi-ai`'s `parseChunkUsage` fills a
 * `cost` object from the model catalog (which does not describe a hand-declared
 * route, so it is all zeros). The raw provider response still carries the real
 * amount, so this plugin reads it there instead.
 *
 * Two observers, joined by generation id:
 *   1. a `globalThis.fetch` wrapper reads `usage.cost` out of provider responses;
 *   2. the `llm/stream` waterfall supplies the session the call belongs to.
 *
 * Neither is enough alone: the wrapper sees money but not the session, the seam
 * sees the session but not the money.
 *
 * @module dsh-session-cost
 */

import fs from 'node:fs';
import os from 'node:os';
import path from 'node:path';

/** Marker so a reloaded plugin does not wrap `fetch` twice. */
const WRAP_MARK = Symbol.for('dsh-session-cost.fetch-wrapper');

/** How long to wait for a cost before calling a completed call unpriced. */
const PENDING_TIMEOUT_MS = 60_000;

/** Upper bounds, so a long-lived host cannot grow without limit. */
const MAX_REGISTRY = 5000;
const MAX_SETTLED = 5000;

/** Where the ledgers live. */
function dataDir() {
  const home = process.env.DSH_HOME || path.join(os.homedir(), '.dsh');
  return path.join(home, 'dsh-session-cost');
}

/** True for the response of a model call we should read. */
function looksLikeModelCall(url, response) {
  if (response.headers.get('x-generation-id')) return true;
  const type = response.headers.get('content-type') || '';
  if (type.includes('text/event-stream')) return true;
  if (!type.includes('json')) return false;
  try {
    const pathname = new URL(url).pathname;
    return /(chat\/completions|completions|messages|responses)$/.test(pathname);
  } catch {
    return false;
  }
}

/** Read a request's URL whichever shape `fetch` was called with. */
function urlOf(input) {
  if (typeof input === 'string') return input;
  if (input instanceof URL) return input.href;
  if (input && typeof input.url === 'string') return input.url;
  return '';
}

/** A finite, non-negative number, or null. */
function finiteAmount(value) {
  if (typeof value !== 'number' || !Number.isFinite(value) || value < 0) return null;
  return value;
}

/** Locate `usage.cost` and the generation id in one parsed provider payload. */
function readPayload(payload) {
  if (!payload || typeof payload !== 'object') return null;
  const usage = payload.usage;
  if (!usage || typeof usage !== 'object') return null;
  const cost = finiteAmount(usage.cost);
  if (cost === null) return null;
  const id = typeof payload.id === 'string' && payload.id.length > 0 ? payload.id : null;
  return { id, cost };
}

/**
 * Split an SSE byte stream into `data:` payloads.
 * Chunk boundaries fall anywhere, so the tail is carried between pieces.
 */
async function eachSsePayload(stream, onPayload) {
  const decoder = new TextDecoder();
  let buffer = '';
  for await (const piece of stream) {
    buffer += decoder.decode(piece, { stream: true });
    const lines = buffer.split('\n');
    buffer = lines.pop() ?? '';
    for (const line of lines) {
      if (!line.startsWith('data:')) continue;
      const body = line.slice(5).trim();
      if (body.length === 0 || body === '[DONE]') continue;
      try {
        onPayload(JSON.parse(body));
      } catch {
        /* a non-JSON keepalive line is not an error */
      }
    }
  }
  buffer += decoder.decode();
  const tail = buffer.trim();
  if (tail.startsWith('data:')) {
    const body = tail.slice(5).trim();
    if (body.length > 0 && body !== '[DONE]') {
      try {
        onPayload(JSON.parse(body));
      } catch {
        /* ignore */
      }
    }
  }
}

export default {
  inject: ['timer'],

  apply(ctx) {
    // ---------------------------------------------------------------- state
    /** generation id -> cost in the provider's own currency. */
    const costByGeneration = new Map();
    /** generation ids already added to a session total. */
    const settled = new Set();
    /** session id -> { total, calls, byPurpose }. */
    const sessionTotals = new Map();
    /** generation id -> { sessionId, purpose }, waiting for its cost. */
    const waiting = new Map();
    /** Diagnostics surfaced on the status page. */
    const stats = {
      observed: 0,        // provider responses whose cost we read
      attributed: 0,      // calls joined to a session
      cancelled: 0,       // aborted by the user; the provider bills nothing
      failed: 0,          // finished with an error; no cost expected
      unpriced: 0,        // completed, but no cost ever arrived
      unattributed: 0,    // cost read, but the call carried no session id
      wrappers: 0,
    };

    let flushTimer = null;
    let dirty = false;

    // ---------------------------------------------------------- persistence
    const stateFile = path.join(dataDir(), 'state.json');

    function load() {
      try {
        const raw = fs.readFileSync(stateFile, 'utf8');
        const parsed = JSON.parse(raw);
        for (const [id, cost] of Object.entries(parsed.costs ?? {})) {
          costByGeneration.set(id, cost);
        }
        for (const [id, total] of Object.entries(parsed.sessions ?? {})) {
          sessionTotals.set(id, {
            total: finiteAmount(total) ?? 0,
            calls: 0,
            byPurpose: {},
          });
        }
        for (const id of parsed.settled ?? []) settled.add(id);
      } catch {
        /* first run, or an unreadable ledger: start clean rather than fail */
      }
    }

    function flush() {
      flushTimer = null;
      if (!dirty) return;
      dirty = false;
      try {
        fs.mkdirSync(path.dirname(stateFile), { recursive: true });
        const payload = {
          version: 1,
          costs: Object.fromEntries(costByGeneration),
          sessions: Object.fromEntries(
            [...sessionTotals].map(([id, value]) => [id, value.total]),
          ),
          settled: [...settled],
        };
        fs.writeFileSync(stateFile, JSON.stringify(payload));
      } catch (error) {
        console.warn('[dsh-session-cost] could not write the ledger:', error?.message ?? error);
      }
    }

    function scheduleFlush() {
      dirty = true;
      if (flushTimer !== null) return;
      flushTimer = ctx.setTimeout(flush, 2000);
    }

    /** Keep the in-memory maps bounded without losing recent history. */
    function trim(map, cap) {
      while (map.size > cap) {
        const oldest = map.keys().next().value;
        map.delete(oldest);
      }
    }

    load();

    // ------------------------------------------------------- cost recording
    /** Add one cost to a session, at most once per generation. */
    function settle(sessionId, generationId, cost, purpose) {
      if (generationId !== null) {
        if (settled.has(generationId)) return;
        settled.add(generationId);
        trim(settled, MAX_SETTLED);
      }
      let entry = sessionTotals.get(sessionId);
      if (entry === undefined) {
        entry = { total: 0, calls: 0, byPurpose: {} };
        sessionTotals.set(sessionId, entry);
      }
      const bucket = purpose === undefined || purpose === '' ? 'agent' : purpose;
      entry.total += cost;
      entry.calls += 1;
      entry.byPurpose[bucket] = (entry.byPurpose[bucket] ?? 0) + cost;
      stats.attributed += 1;
      scheduleFlush();
    }

    /** A cost arrived: join it to a waiting call, or park it for a later one. */
    function recordCost(generationId, cost) {
      stats.observed += 1;
      const parked = waiting.get(generationId);
      if (parked !== undefined) {
        waiting.delete(generationId);
        settle(parked.sessionId, generationId, cost, parked.purpose);
        return;
      }
      if (generationId === null) {
        // No id to join on: the money is real but cannot be placed.
        stats.unattributed += 1;
        console.warn('[dsh-session-cost] a provider response reported cost without an id');
        return;
      }
      costByGeneration.set(generationId, cost);
      trim(costByGeneration, MAX_REGISTRY);
      scheduleFlush();
      // If nothing ever comes back for it, it stays here and is reported as
      // parked rather than silently dropped.
    }

    // ---------------------------------------------------- observer 1: fetch
    function installFetchWrapper() {
      const current = globalThis.fetch;
      if (typeof current !== 'function') return;
      if (current[WRAP_MARK] === true) {
        stats.wrappers = 1;
        return; // a previous instance already wrapped it
      }

      const wrapped = async function (input, init) {
        const response = await current.call(this, input, init);
        try {
          const url = urlOf(input);
          if (!looksLikeModelCall(url, response)) return response;
          'use strict';
          const body = response.body;
          if (body === null || typeof body.tee !== 'function') return response;

          const generationId =
            response.headers.get('x-generation-id') ??
            (typeof response.headers.get('x-session-id') === 'string'
              ? response.headers.get('x-session-id')
              : null);

          const [mine, theirs] = body.tee();
          const type = response.headers.get('content-type') || '';

          // Read our branch to completion; never leave it unread, or the
          // consumer's branch would be back-pressured behind it.
          const observe = (async () => {
            try {
              if (type.includes('text/event-stream')) {
                let id = generationId;
                await eachSsePayload(mine, (payload) => {
                  const found = readPayload(payload);
                  if (found === null) return;
                  if (found.id !== null) id = found.id;
                  if (id !== null) recordCost(id, found.cost);
                });
              } else {
                const text = await new Response(mine).text();
                let parsed = null;
                try {
                  parsed = JSON.parse(text);
                } catch {
                  return;
                }
                const found = readPayload(parsed);
                if (found !== null) {
                  const id = found.id ?? generationId;
                  if (id !== null) recordCost(id, found.cost);
                }
              }
            } catch {
              /* an aborted or truncated body carries no cost; nothing to do */
            }
          })();
          // Detach: the consumer must not wait on our bookkeeping.
          void observe;

          const rebuilt = new Response(theirs, {
            status: response.status,
            statusText: response.statusText,
            headers: response.headers,
          });
          // A rebuilt Response loses `url`; the OpenAI SDK logs it on retries.
          try {
            Object.defineProperty(rebuilt, 'url', { value: response.url, configurable: true });
          } catch {
            /* not fatal: url is used for diagnostics only */
          }
          stats.wrappers = 1;
          return rebuilt;
        } catch (error) {
          // Any failure here must leave the request working.
          console.warn('[dsh-session-cost] read failed, passing the response through:', error?.message ?? error);
          return response;
        }
      };

      Object.defineProperty(wrapped, WRAP_MARK, { value: true });
      globalThis.fetch = wrapped;
      ctx.effect(() => () => {
        if (globalThis.fetch === wrapped) globalThis.fetch = current;
      }, 'dsh-session-cost: fetch wrapper');
    }

    installFetchWrapper();

    // ----------------------------------------------------- observer 2: seam
    ctx.on('llm/stream', (options, next) => {
      const sessionId = typeof options?.sessionId === 'string' ? options.sessionId : undefined;
      const purpose = typeof options?.purpose === 'string' ? options.purpose : undefined;
      return (async function* passThrough() {
        let generationId = null;
        let finished = null;
        try {
          for await (const chunk of await next()) {
            if (chunk && chunk.type === 'usage') {
              // Nothing to take: the harness's usage has no cost field.
            } else if (chunk && chunk.type === 'finish') {
              finished = chunk.reason ?? null;
              const id = chunk.replayState?.response?.responseId;
              if (typeof id === 'string' && id.length > 0) generationId = id;
            }
            yield chunk;
          }
        } finally {
          const kind = finished?.kind ?? 'stop';
          if (kind === 'aborted') {
            // The provider bills nothing for a cancelled generation (verified:
            // a late abort still reports total_cost 0), so this is not a loss.
            stats.cancelled += 1;
          } else if (kind === 'error') {
            stats.failed += 1;
          } else if (generationId === null) {
            stats.unpriced += 1;
          } else if (costByGeneration.has(generationId)) {
            const cost = costByGeneration.get(generationId);
            costByGeneration.delete(generationId);
            if (sessionId === undefined) {
              stats.unattributed += 1;
            } else {
              settle(sessionId, generationId, cost, purpose);
            }
          } else if (sessionId === undefined) {
            // Future plugins (e.g. an experimental auto-review) call the model
            // with no sessionId at all; those calls cannot be placed.
            stats.unattributed += 1;
          } else {
            // Normal path is cost-first, so this is rare: park the call and let
            // the cost find it.
            waiting.set(generationId, { sessionId, purpose, at: Date.now() });
          }
        }
      })();
    });

    // Drop waiters whose cost never arrived, and report them.
    ctx.setInterval(() => {
      const now = Date.now();
      for (const [id, entry] of waiting) {
        if (now - entry.at < PENDING_TIMEOUT_MS) continue;
        waiting.delete(id);
        stats.unpriced += 1;
        console.warn('[dsh-session-cost] a completed call never reported a cost:', id);
      }
    }, 30_000);

    // --------------------------------------------------------------- status
    function snapshot(sessionId) {
      const entry = sessionTotals.get(sessionId);
      return {
        sessionId,
        total: entry?.total ?? null,
        calls: entry?.calls ?? 0,
        byPurpose: entry?.byPurpose ?? {},
        currency: 'RUB',
        stats: { ...stats, waiting: waiting.size, parked: costByGeneration.size, sessions: sessionTotals.size },
      };
    }

    // ---------------------------------------------------------------- route
    const ROUTE_PREFIX = '/_dsh/dsh-session-cost';

    /** Answer with JSON, or with a plain error the client can show. */
    function respond(res, status, payload) {
      const body = JSON.stringify(payload);
      res.writeHead(status, {
        'content-type': 'application/json; charset=utf-8',
        'cache-control': 'no-store',
        'content-length': Buffer.byteLength(body),
      });
      res.end(body);
    }

    ctx.inject(['connection', 'webServer'], (webCtx) => {
      webCtx.effect(() => {
        const dispose = webCtx.webServer.register({
          kind: 'prefix',
          path: ROUTE_PREFIX,
          handler: async (req, res) => {
            try {
              const rejection = webCtx.connection.requestRejection(req);
              if (rejection !== undefined) {
                respond(res, rejection, { error: rejection === 401 ? 'authentication required' : 'request rejected' });
                return;
              }
              const url = new URL(req.url || '/', 'http://localhost');
              const tail = url.pathname.slice(ROUTE_PREFIX.length);
              if (tail !== '/session' && tail !== '/status') {
                respond(res, 404, { error: 'unknown method' });
                return;
              }
              const sessionId = url.searchParams.get('sessionId') ?? '';
              respond(res, 200, tail === '/session' ? snapshot(sessionId) : snapshot(''));
            } catch (error) {
              respond(res, 500, { error: String(error?.message ?? error) });
            }
          },
        });
        return dispose;
      }, 'dsh-session-cost: http route');
    });

    ctx.effect(() => () => {
      if (flushTimer !== null) ctx.clearTimeout(flushTimer);
      flush();
    }, 'dsh-session-cost: ledger flush');
  },
};