import { LatencyTracker } from './LatencyTracker.js';
import { LLMResponse } from './types.js';
import { CostCalculator } from './CostCalculator.js';

/**
 * Logger interface for ModelTimeout
 */
export interface ModelTimeoutLogger {
  info: (message: string, meta?: Record<string, unknown>) => void;
  warn: (message: string, meta?: Record<string, unknown>) => void;
  error: (message: string, meta?: Record<string, unknown>) => void;
}

// Default console logger
const consoleLogger: ModelTimeoutLogger = {
  info: (message: string, meta?: Record<string, unknown>) =>
    meta ? console.info(message, meta) : console.info(message),
  warn: (message: string, meta?: Record<string, unknown>) =>
    meta ? console.warn(message, meta) : console.warn(message),
  error: (message: string, meta?: Record<string, unknown>) =>
    meta ? console.error(message, meta) : console.error(message)
};

let modelTimeoutLogger: ModelTimeoutLogger = consoleLogger;

export const setModelTimeoutLogger = (logger: ModelTimeoutLogger): void => {
  modelTimeoutLogger = logger;
};

export class ModelTimeoutError extends Error {
  constructor(public readonly model: string, public readonly timeoutMs: number) {
    super(`Model "${model}" timed out after ${timeoutMs}ms`);
    this.name = 'ModelTimeoutError';
  }
}

/**
 * Wrapper to add timeout to LLM calls with latency tracking
 */
export async function callWithTimeout(
  model: string,
  call: (signal?: AbortSignal) => Promise<LLMResponse>,
  timeoutMs: number,
  latencyTracker?: LatencyTracker,
  stage?: string
): Promise<LLMResponse> {
  const abortController = new AbortController();
  let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
  let isResolved = false;
  const timeoutError = new ModelTimeoutError(model, timeoutMs);
  const invocationStart = Date.now();

  const callPromise = call(abortController.signal);
  const timeoutPromise = new Promise<never>((_, reject) => {
    timeoutHandle = setTimeout(() => {
      if (!isResolved) {
        abortController.abort();
        reject(timeoutError);
      }
    }, timeoutMs);
  });

  try {
    const response = await Promise.race([callPromise, timeoutPromise]);
    isResolved = true;

    const durationMs = Date.now() - invocationStart;
    const responseModel = response.model ?? model;
    const tokens = response.tokens ?? null;
    const cost = normalizeCost(response.cost, responseModel);
    latencyTracker?.recordModelLatency(model, durationMs, 'success', { stage, timeoutMs, tokens, cost });
    modelTimeoutLogger.info('Model call succeeded', { model, durationMs, stage, timeoutMs, tokens, cost });

    if (!response.model) {
      response.model = model;
    }
    return response;
  } catch (error) {
    isResolved = true;

    const durationMs = Date.now() - invocationStart;
    const message = error instanceof Error ? error.message : 'Unknown error';
    latencyTracker?.recordModelLatency(model, durationMs, 'error', {
      stage,
      timeoutMs,
      errorMessage: message,
      tokens: null,
      cost: null
    });

    if (error instanceof ModelTimeoutError || abortController.signal.aborted) {
      // Don't wait for the aborted promise to prevent hanging
      callPromise.catch((callError) => {
        if (callError instanceof Error && callError.name !== 'AbortError') {
          modelTimeoutLogger.warn(`Model ${model} aborted with non-standard error`, { error: callError.message });
        }
      });
      modelTimeoutLogger.warn(`Model ${model} timed out after ${timeoutMs}ms`, { durationMs, stage });
      return buildTimeoutFallback(model);
    }

    modelTimeoutLogger.error(`Model ${model} call failed`, { error: message, durationMs, stage });
    return buildModelErrorResponse(model, message);
  } finally {
    // Always clear timeout to prevent memory leaks
    if (timeoutHandle !== undefined) {
      clearTimeout(timeoutHandle);
      timeoutHandle = undefined;
    }
  }
}

/**
 * Build timeout fallback response
 */
export function buildTimeoutFallback(model: string): LLMResponse & { error: true } {
  return {
    response: 'Timeout: Response took too long',
    tokens: 0,
    cost: 0,
    model,
    error: true
  };
}

/**
 * Build model error response
 */
export function buildModelErrorResponse(model: string, message: string): LLMResponse & { error: true } {
  return {
    response: `Error: ${message}`,
    tokens: 0,
    cost: 0,
    model,
    error: true
  };
}

/**
 * Normalize cost to valid number or null
 */
function normalizeCost(cost: number | null | undefined, context?: string): number | null {
  if (cost === null || cost === undefined) {
    return null;
  }

  return CostCalculator.validateCost(cost, context ?? 'unknown');
}

/**
 * Run model call with optional delay before execution
 */
export async function runModelCallWithDelay(
  delayMs: number,
  model: string,
  call: (signal?: AbortSignal) => Promise<LLMResponse>,
  timeoutMs: number,
  latencyTracker?: LatencyTracker,
  stage?: string
): Promise<LLMResponse> {
  const startTime = Date.now();
  modelTimeoutLogger.info(`[TIMING] Starting model call for ${model}`, { model, delayMs, timeoutMs });

  if (delayMs > 0) {
    await new Promise<void>((resolve) => setTimeout(resolve, delayMs));
    modelTimeoutLogger.info(`[TIMING] Delay complete for ${model}`, { model, delayMs, elapsed: Date.now() - startTime });
  }

  const beforeCallTime = Date.now();
  const result = await callWithTimeout(model, call, timeoutMs, latencyTracker, stage);
  const callDuration = Date.now() - beforeCallTime;
  const totalDuration = Date.now() - startTime;
  const validatedCost = normalizeCost(result.cost, result.model ?? model);

  modelTimeoutLogger.info(`[TIMING] Model call complete for ${model}`, {
    model,
    stage,
    callDuration,
    totalDuration,
    tokens: result.tokens ?? null,
    cost: validatedCost
  });

  return result;
}