import { Axiom } from '@axiomhq/js';
import { logger } from './color';
import { getCorrelationId, requestToContext } from './common';
import type {
  BaseLog,
  ChildLogOptions,
  ConsoleLog,
  Log,
  LogLevel,
} from './type';

let lastMessageTimestamp: number | undefined;

export function createBaseLog(
  {
    token,
    orgId = 'vrplatform-mv6k',
    dataset,
    consoleLog,
  }: {
    token?: string;
    orgId?: string;
    dataset?: string;
    consoleLog?: ConsoleLog;
  },
  {
    environment,
    workerId,
    app,
    version,
    context,
    type,
    correlationId,
    sessionId,
    executionContext,
  }: {
    correlationId?: string;
    sessionId?: string;
    environment?: 'development' | 'production';
    type?: 'worker' | 'durableObject';
    workerId?: string;
    app?: string;
    version?: string;
    context?: Record<string, any>;
    executionContext?: {
      waitUntil(promise: Promise<any>): void;
    };
  }
): BaseLog {
  const baseFields: Record<string, any> = {
    type: type || 'worker',
    environment: environment || 'production',
    worker: {
      id: workerId || 'default',
      app: app || 'default',
      version: version || 'default',
      started: new Date().toISOString(),
    },
    ...context,
  };

  const axiom = token
    ? new Axiom({
        token,
        orgId,
        onError(err) {
          console.log('Axiom error', err);
        },
      })
    : undefined;

  function ingest(
    level: LogLevel,
    message: string,
    additionalFields: Record<string, any> = {}
  ) {
    axiom?.ingest(dataset!, {
      _time: Date.now(),
      level,
      message,
      ...baseFields,
      ...additionalFields,
    });
  }

  const sid = sessionId || crypto.randomUUID().substring(0, 13);

  function createBaseLog(arg: ChildLogOptions = {}): Log {
    const childContext = arg.context || {};
    if (arg.request) {
      childContext.request = requestToContext(arg.request);
    }
    const baseLog = (
      level: LogLevel,
      message: any,
      extra: Record<string, any> = {},
      error?: Error
    ) => {
      const cid =
        arg.correlationId ||
        correlationId ||
        childContext?.correlationId ||
        extra?.correlationId ||
        correlationId;
      ingest(level, message, {
        ...childContext,
        ...extra,
        correlationId: cid,
        sessionId: sid,
      });
      if (consoleLog) {
        const name = arg.name ? `${app}/${arg.name}` : app;
        const lev = level.toUpperCase();

        const now = Date.now();
        const diff = (
          lastMessageTimestamp ? `${now - lastMessageTimestamp}ms ` : ''
        ).padEnd(6, ' ');
        lastMessageTimestamp = now;

        error = message instanceof Error ? message : error;
        message = message instanceof Error ? message.message : message;
        const s = [sid, cid].filter(Boolean).join('/');
        const formatted = `[${lev}] [${name}]${s ? ` [${s}]` : ''} ${message}`;
        if (typeof consoleLog === 'function')
          consoleLog({
            formatted,
            level,
            message,
            extra,
            error,
            context: childContext,
            tags: {
              sessionId: sid,
              correlationId: cid,
              environment,
              workerId,
              app,
              version,
              type,
              name,
            },
          });
        else if (consoleLog === 'color')
          logger
            .dim()
            .append(`${diff}`)
            .reset()
            .color('green')
            .append(`[${name}] `)
            .reset()
            .bold()
            .append(`${lev} `)
            .reset()
            .color('blue')
            .bold()
            .append(s ? `[${s}] ` : '')
            .reset()
            .append(message)
            .log();
        else console.log(formatted, extra);

        if (error) console.error(error);
      }
    };
    return {
      addContext(arg) {
        for (const key in arg) {
          childContext[key] = arg[key];
        }
      },
      child: ({
        correlationId,
        name,
        request,
        context = {},
      }: ChildLogOptions) =>
        createBaseLog({
          correlationId:
            correlationId ||
            (request ? getCorrelationId(request) : undefined) ||
            undefined,
          name,
          request,
          context: { ...childContext, ...context },
        }),
      error: (message, error, data) => {
        if (message instanceof Error) {
          return baseLog('error', message.message, error as any, message);
        }
        const m =
          typeof message === 'string'
            ? message
            : (error as any)?.message || (error as any)?.name || 'Error';
        if (error && error instanceof Error)
          return baseLog('error', m, data, error);
        return baseLog(
          'error',
          m,
          error && typeof error === 'object' ? error : data
        );
      },
      debug: (message, data) => baseLog('debug', message, data),
      warn: (message, data) => baseLog('warn', message, data),
      info: (message, data) => baseLog('info', message, data),
    };
  }
  const baseLog = createBaseLog() as any as BaseLog;
  baseLog.flush = () => {
    if (executionContext && axiom) {
      executionContext.waitUntil(axiom.flush());
    } else if (axiom) return axiom.flush();
  };
  return baseLog;
}