import { mkdir, writeFile } from 'node:fs/promises';
import path from 'node:path';
import { execSync } from 'node:child_process';

export type LatencyLogger = {
  info: (message: string, meta?: Record<string, unknown>) => unknown;
  warn: (message: string, meta?: Record<string, unknown>) => unknown;
  error: (message: string, meta?: Record<string, unknown>) => unknown;
};

const consoleLogger: LatencyLogger = {
  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 latencyLogger: LatencyLogger = consoleLogger;

export const setLatencyTrackerLogger = (logger: LatencyLogger): void => {
  latencyLogger = logger;
};

const resolveVerboseFlag = (value: string | undefined | null): boolean => {
  if (value === undefined || value === null) {
    return true;
  }

  const normalized = value.trim().toLowerCase();
  return !['false', '0', 'no', 'off', 'disabled'].includes(normalized);
};

const verboseLoggingEnv = process.env.MCP_VERBOSE_LATENCY_LOGGING;
const legacyVerboseLoggingEnv = process.env.SECOND_OPINION_VERBOSE_LOGGING;

let verboseLatencyLogging = resolveVerboseFlag(verboseLoggingEnv ?? legacyVerboseLoggingEnv);

export const setVerboseLatencyLogging = (enabled: boolean): void => {
  verboseLatencyLogging = enabled;
};

export const isVerboseLatencyLoggingEnabled = (): boolean => verboseLatencyLogging;

export type StageMetric = {
  name: string;
  durationMs: number;
  meta?: Record<string, unknown>;
};

export type ModelMetric = {
  model: string;
  durationMs: number;
  status: 'success' | 'error';
  stage?: string;
  errorMessage?: string;
  timeoutMs?: number;
  tokens?: number | null;
  cost?: number | null;
};

export type LatencySummaryTotals = {
  totalTokens?: number;
  totalCost?: number;
  totalModels?: number;
};

type PersistMeta = {
  errorMessage?: string;
  latencySummary?: LatencySummaryTotals;
} & Record<string, unknown>;

type SummaryOptions = {
  message?: string;
  totals?: LatencySummaryTotals;
  errorMessage?: string;
  includeContextKeys?: string[];
};

export class LatencyTracker {
  private static metricsDirectory: string | null = null;

  private readonly overallStart = Date.now();
  private readonly stageTimers = new Map<string, number>();
  private readonly stageMetrics: StageMetric[] = [];
  private readonly modelMetrics: ModelMetric[] = [];
  private readonly context: Record<string, unknown> = {};

  constructor(private readonly flowName: string, private readonly baseMeta: Record<string, unknown> = {}) {}

  startStage(stage: string): void {
    this.stageTimers.set(stage, Date.now());
  }

  isStageActive(stage: string): boolean {
    return this.stageTimers.has(stage);
  }

  endStage(stage: string, meta: Record<string, unknown> = {}): number {
    const start = this.stageTimers.get(stage);
    if (!start) {
      if (verboseLatencyLogging) {
        latencyLogger.warn('Attempted to end inactive stage', { flow: this.flowName, stage });
      }
      return 0;
    }

    const durationMs = Date.now() - start;
    this.stageMetrics.push({ name: stage, durationMs, meta });
    this.stageTimers.delete(stage);
    if (verboseLatencyLogging) {
      latencyLogger.info('Stage complete', { flow: this.flowName, stage, durationMs, ...meta });
    }
    return durationMs;
  }

  recordStage(stage: string, durationMs: number, meta: Record<string, unknown> = {}): void {
    this.stageMetrics.push({ name: stage, durationMs, meta });
    if (verboseLatencyLogging) {
      latencyLogger.info('Stage recorded', { flow: this.flowName, stage, durationMs, ...meta });
    }
  }

  recordModelLatency(
    model: string,
    durationMs: number,
    status: 'success' | 'error',
    meta: { stage?: string; errorMessage?: string; timeoutMs?: number; tokens?: number | null; cost?: number | null } = {}
  ): void {
    this.modelMetrics.push({
      model,
      durationMs,
      status,
      stage: meta.stage,
      errorMessage: meta.errorMessage,
      timeoutMs: meta.timeoutMs,
      tokens: meta.tokens ?? null,
      cost: meta.cost ?? null
    });

    const logPayload = {
      flow: this.flowName,
      model,
      durationMs,
      status,
      stage: meta.stage,
      timeoutMs: meta.timeoutMs,
      error: meta.errorMessage,
      tokens: meta.tokens ?? null,
      cost: meta.cost ?? null
    };

    if (verboseLatencyLogging) {
      if (status === 'success') {
        latencyLogger.info('Model latency recorded', logPayload);
      } else {
        latencyLogger.warn('Model latency recorded with error', logPayload);
      }
    }
  }

  addContext(meta: Record<string, unknown>): void {
    Object.assign(this.context, meta);
  }

  endAllActiveStages(errorMessage?: string): void {
    const activeStages = Array.from(this.stageTimers.keys());
    for (const stage of activeStages) {
      this.endStage(stage, { error: errorMessage });
    }
  }

  getSnapshot(): {
    totalDurationMs: number;
    stages: StageMetric[];
    models: ModelMetric[];
    context: Record<string, unknown>;
  } {
    const totalDurationMs = Date.now() - this.overallStart;
    return {
      totalDurationMs,
      stages: this.stageMetrics.map((stage) => ({
        name: stage.name,
        durationMs: stage.durationMs,
        meta: stage.meta ? { ...stage.meta } : undefined
      })),
      models: this.modelMetrics.map((model) => ({
        model: model.model,
        durationMs: model.durationMs,
        status: model.status,
        stage: model.stage,
        errorMessage: model.errorMessage,
        timeoutMs: model.timeoutMs,
        tokens: model.tokens ?? null,
        cost: model.cost ?? null
      })),
      context: { ...this.baseMeta, ...this.context }
    };
  }

  buildSummaryLog(status: 'success' | 'error', options: SummaryOptions = {}): Record<string, unknown> {
    const snapshot = this.getSnapshot();

    const longestStage = snapshot.stages.reduce<(typeof snapshot.stages[number]) | null>(
      (current, candidate) => {
        if (!current || candidate.durationMs > current.durationMs) {
          return candidate;
        }
        return current;
      },
      null
    );

    const slowestModel = snapshot.models.reduce<(typeof snapshot.models[number]) | null>(
      (current, candidate) => {
        if (!current || candidate.durationMs > current.durationMs) {
          return candidate;
        }
        return current;
      },
      null
    );

    const summary: Record<string, unknown> = {
      flow: this.flowName,
      status,
      totalDurationMs: snapshot.totalDurationMs,
      totalStages: snapshot.stages.length,
      totalModelCalls: snapshot.models.length
    };

    if (options.includeContextKeys && options.includeContextKeys.length > 0) {
      const allowedEntries = options.includeContextKeys
        .filter((key) => Object.prototype.hasOwnProperty.call(snapshot.context, key))
        .map((key) => [key, snapshot.context[key]] as const);
      if (allowedEntries.length > 0) {
        summary.context = Object.fromEntries(allowedEntries);
      }
    }

    if (longestStage) {
      summary.longestStage = longestStage.name;
      summary.longestStageDurationMs = longestStage.durationMs;
      if (longestStage.meta) {
        summary.longestStageMeta = { ...longestStage.meta };
      }
    }

    if (slowestModel) {
      summary.slowestModel = slowestModel.model;
      summary.slowestModelDurationMs = slowestModel.durationMs;
      summary.slowestModelStage = slowestModel.stage ?? null;
      summary.slowestModelStatus = slowestModel.status;
      if (slowestModel.tokens != null) {
        summary.slowestModelTokens = slowestModel.tokens;
      }
      if (slowestModel.cost != null) {
        summary.slowestModelCost = slowestModel.cost;
      }
      if (slowestModel.errorMessage) {
        summary.slowestModelError = slowestModel.errorMessage;
      }
    }

    if (options.totals) {
      summary.responseTotals = options.totals;
    }

    if (options.errorMessage) {
      summary.errorMessage = options.errorMessage;
    }

    return summary;
  }

  logSummary(status: 'success' | 'error', options: SummaryOptions = {}): Record<string, unknown> | null {
    if (!verboseLatencyLogging) {
      return null;
    }

    const summary = this.buildSummaryLog(status, options);
    const message = options.message ?? 'Latency summary';
    latencyLogger.info(message, summary);
    return summary;
  }

  async persist(status: 'success' | 'error', meta: PersistMeta = {}): Promise<void> {
    const totalDurationMs = Date.now() - this.overallStart;
    if (process.env.JEST_WORKER_ID) {
      return;
    }
    const payload = {
      flow: this.flowName,
      status,
      totalDurationMs,
      timestamp: new Date().toISOString(),
      stages: this.stageMetrics,
      models: this.modelMetrics,
      context: { ...this.baseMeta, ...this.context },
      meta
    };

    try {
      const metricsDir = await LatencyTracker.ensureMetricsDirectory();
      const filePath = path.join(metricsDir, `${this.flowName}-${Date.now()}.json`);
      await writeFile(filePath, JSON.stringify(payload, null, 2), { encoding: 'utf-8' });
      if (verboseLatencyLogging) {
        latencyLogger.info('Latency metrics persisted', { flow: this.flowName, status, totalDurationMs, filePath });
      }
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : 'Unknown error writing latency metrics';
      latencyLogger.error('Failed to persist latency metrics', { flow: this.flowName, status, error: errorMessage });
    }
  }
  private static async ensureMetricsDirectory(): Promise<string> {
    if (this.metricsDirectory) {
      return this.metricsDirectory;
    }

    // Use process.cwd() for Jest compatibility (import.meta.url breaks Jest transforms)
    const repoRoot = process.cwd();
    const repoName = path.basename(repoRoot);

    let branch = process.env.GIT_BRANCH || process.env.BRANCH_NAME || process.env.CURRENT_BRANCH || '';
    if (!branch || branch === 'HEAD') {
      try {
        // Validate repoRoot path to prevent command injection
        const normalizedPath = path.normalize(repoRoot);
        if (normalizedPath.includes('..') || !path.isAbsolute(normalizedPath)) {
          throw new Error('Invalid repository path detected');
        }

        branch = execSync('git rev-parse --abbrev-ref HEAD', {
          cwd: normalizedPath,
          stdio: ['ignore', 'pipe', 'ignore'],
          timeout: 5000 // Add timeout for safety
        }).toString().trim();
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : 'Unknown error resolving branch';
        latencyLogger.warn('Unable to resolve git branch for latency metrics, defaulting to "unknown"', { error: errorMessage });
        branch = 'unknown';
      }
    }

    // Sanitize branch name for filesystem (replace slashes with hyphens)
    const sanitizedBranch = (branch || 'unknown').replace(/\//g, '-');
    const metricsDir = path.join('/tmp', repoName, sanitizedBranch);

    await mkdir(metricsDir, { recursive: true });
    this.metricsDirectory = metricsDir;
    return metricsDir;
  }
}