/**
 * AI Generations Service
 *
 * Query and format AiContext records for the AI Processing dashboard.
 * Provides paginated, filterable listings with aggregate statistics.
 */

import { getModels } from '../../models';
import { escapeRegex } from '../../utils/escapeRegex';

// ============================================
// COST ESTIMATION (per 1M tokens, USD)
// ============================================

const MODEL_PRICING: Record<string, { inputPerMillion: number; outputPerMillion: number }> = {
  // Anthropic Claude
  'claude-3-5-sonnet': { inputPerMillion: 3, outputPerMillion: 15 },
  'claude-3-5-sonnet-20241022': { inputPerMillion: 3, outputPerMillion: 15 },
  'claude-3-5-sonnet-20240620': { inputPerMillion: 3, outputPerMillion: 15 },
  'claude-3-opus': { inputPerMillion: 15, outputPerMillion: 75 },
  'claude-3-opus-20240229': { inputPerMillion: 15, outputPerMillion: 75 },
  'claude-3-haiku': { inputPerMillion: 0.25, outputPerMillion: 1.25 },
  'claude-3-haiku-20240307': { inputPerMillion: 0.25, outputPerMillion: 1.25 },
  'claude-sonnet-4-20250514': { inputPerMillion: 3, outputPerMillion: 15 },
  'claude-opus-4-20250514': { inputPerMillion: 15, outputPerMillion: 75 },
  // OpenAI
  'gpt-4o': { inputPerMillion: 2.5, outputPerMillion: 10 },
  'gpt-4o-mini': { inputPerMillion: 0.15, outputPerMillion: 0.6 },
  'gpt-4-turbo': { inputPerMillion: 10, outputPerMillion: 30 },
  'gpt-4': { inputPerMillion: 30, outputPerMillion: 60 },
  'gpt-3.5-turbo': { inputPerMillion: 0.5, outputPerMillion: 1.5 },
  // Zhipu GLM (also served locally via the ollama provider).
  // 'glm-5' covers glm-5, glm-5.1 and glm-5:cloud via prefix matching.
  'glm-4': { inputPerMillion: 1.4, outputPerMillion: 1.4 },
  'glm-4-plus': { inputPerMillion: 3.5, outputPerMillion: 3.5 },
  'glm-4-flash': { inputPerMillion: 0.1, outputPerMillion: 0.1 },
  'glm-5': { inputPerMillion: 0.6, outputPerMillion: 2.2 },
};

// Fallback pricing for any model not explicitly listed above (incl. "unknown"
// or a missing model). Ensures every record with token usage gets an estimate
// instead of a blank cost. Mid-range blended estimate.
const DEFAULT_PRICING = { inputPerMillion: 0.6, outputPerMillion: 2.2 };

function resolvePricing(aiModel?: string | null): { inputPerMillion: number; outputPerMillion: number } {
  if (aiModel) {
    if (MODEL_PRICING[aiModel]) return MODEL_PRICING[aiModel];
    const prefix = Object.keys(MODEL_PRICING).find((k) => aiModel.startsWith(k));
    if (prefix) return MODEL_PRICING[prefix];
  }
  return DEFAULT_PRICING;
}

/**
 * Estimate total cost across all records for a company.
 * Groups by aiModel, sums input/output tokens, applies per-model pricing.
 * Unknown models fall back to DEFAULT_PRICING, and records that only recorded a
 * combined `tokensUsed` total (no input/output split) are estimated from it.
 */
async function estimateTotalCost(companyId: string): Promise<number | null> {
  const { AiContext } = getModels();

  const tokenByModel = await AiContext.aggregate([
    { $match: { companyId } },
    {
      $group: {
        _id: '$aiModel',
        totalInput: { $sum: { $ifNull: ['$inputTokens', 0] } },
        totalOutput: { $sum: { $ifNull: ['$outputTokens', 0] } },
        // Combined total for records lacking an input/output split.
        splitlessTotal: {
          $sum: {
            $cond: [
              { $or: [{ $gt: ['$inputTokens', 0] }, { $gt: ['$outputTokens', 0] }] },
              0,
              { $ifNull: ['$tokensUsed', 0] },
            ],
          },
        },
      },
    },
  ]);

  if (!tokenByModel.length) return null;

  let totalCost = 0;
  let hasCost = false;

  for (const row of tokenByModel) {
    const pricing = resolvePricing(row._id as string);
    if (row.totalInput > 0 || row.totalOutput > 0) {
      totalCost += (row.totalInput / 1_000_000) * pricing.inputPerMillion;
      totalCost += (row.totalOutput / 1_000_000) * pricing.outputPerMillion;
      hasCost = true;
    }
    if (row.splitlessTotal > 0) {
      const blended = (pricing.inputPerMillion + pricing.outputPerMillion) / 2;
      totalCost += (row.splitlessTotal / 1_000_000) * blended;
      hasCost = true;
    }
  }

  return hasCost ? totalCost : null;
}

/**
 * Estimate cost for a single record based on its model and token counts.
 * Uses the exact input/output split when available, otherwise falls back to the
 * combined total. Returns null only when there is no token data at all.
 */
function estimateRecordCost(
  aiModel: string | null,
  inputTokens: number | null,
  outputTokens: number | null,
  totalTokens?: number | null,
): number | null {
  const pricing = resolvePricing(aiModel);

  if (inputTokens != null && outputTokens != null) {
    return (inputTokens / 1_000_000) * pricing.inputPerMillion + (outputTokens / 1_000_000) * pricing.outputPerMillion;
  }

  const total = totalTokens ?? inputTokens ?? outputTokens ?? null;
  if (total == null) return null;
  const blended = (pricing.inputPerMillion + pricing.outputPerMillion) / 2;
  return (total / 1_000_000) * blended;
}

// ============================================
// TYPES
// ============================================

export interface AiGenerationRecord {
  id: string;
  companyId: string;
  moduleSource: string;
  entityId?: string | null;
  analysisType: string;
  aiModel: string;
  provider: string;
  tokensUsed: number | null;
  inputTokens: number | null;
  outputTokens: number | null;
  processingTimeMs: number | null;
  latencyMs: number | null;
  overallConfidence: number | null;
  finishReason: string | null;
  apiKeyMasked: string | null;
  status: string;
  createdAt: string;
  updatedAt: string;
  completedAt: string | null;
  estimatedCost: number | null;
  generatedFields: string[];
  inputSummary: string;
  pipelineVersion: string;
  // Error info for failed contexts (mapped from low confidence or rejected status)
  error?: string | null;
}

export interface AiGenerationDetailRecord extends AiGenerationRecord {
  analysis: Record<string, any> | null;
}

export interface AiGenerationStats {
  totalGenerations: number;
  completedCount: number;
  failedCount: number;
  processingCount: number;
  avgProcessingTimeMs: number | null;
  /** Sum of processingTimeMs across the filtered set — the Total Duration card. */
  totalProcessingTimeMs: number;
  totalInputTokens: number;
  totalOutputTokens: number;
  totalTokensUsed: number;
  estimatedCost: number | null;
  providerBreakdown: Record<string, number>;
  moduleBreakdown: Record<string, number>;
}

export interface AiGenerationListResponse {
  records: AiGenerationRecord[];
  total: number;
  page: number;
  limit: number;
  totalPages: number;
  stats: AiGenerationStats;
}

// ============================================
// HELPER: Transform AiContext doc to flat record
// ============================================

function formatRecord(doc: any): AiGenerationRecord {
  const analysis = doc.analysis || {};
  const inputs = doc.inputs || {};

  // Build a summary of generated fields (top-level keys that have values)
  const generatedFields = Object.keys(analysis).filter(
    (key) => analysis[key] !== undefined && analysis[key] !== null && analysis[key] !== ''
  );

  // Build input summary (company name, website, or raw text preview)
  let inputSummary = '';
  if (inputs.companyName) {
    inputSummary = inputs.companyName;
  } else if (inputs.rawText) {
    inputSummary = inputs.rawText.substring(0, 150) + (inputs.rawText.length > 150 ? '...' : '');
  } else if (inputs.description) {
    inputSummary = inputs.description.substring(0, 150) + (inputs.description.length > 150 ? '...' : '');
  }

  // Map status: AiContext uses draft/reviewed/approved/rejected
  // Map to dashboard-friendly statuses
  let displayStatus = doc.status || 'draft';
  if (displayStatus === 'rejected') displayStatus = 'failed';
  else if (displayStatus === 'reviewed' || displayStatus === 'approved') displayStatus = 'completed';
  else if (displayStatus === 'draft') displayStatus = 'completed'; // drafts are still completed generations

  return {
    id: doc._id?.toString?.() || doc.id?.toString?.() || '',
    companyId: doc.companyId || '',
    moduleSource: doc.moduleSource || '',
    entityId: doc.entityId || null,
    analysisType: doc.analysisType || '',
    aiModel: doc.aiModel || '',
    provider: doc.provider || '',
    tokensUsed: doc.tokensUsed ?? null,
    inputTokens: doc.inputTokens ?? null,
    outputTokens: doc.outputTokens ?? null,
    processingTimeMs: doc.processingTimeMs ?? null,
    latencyMs: doc.latencyMs ?? null,
    overallConfidence: doc.overallConfidence ?? null,
    finishReason: doc.finishReason ?? null,
    apiKeyMasked: doc.apiKeyMasked ?? null,
    status: displayStatus,
    createdAt: doc.createdAt?.toISOString?.() || new Date().toISOString(),
    updatedAt: doc.updatedAt?.toISOString?.() || new Date().toISOString(),
    completedAt: doc.completedAt?.toISOString?.() || null,
    estimatedCost: estimateRecordCost(doc.aiModel, doc.inputTokens, doc.outputTokens, doc.tokensUsed),
    generatedFields,
    inputSummary,
    pipelineVersion: doc.pipelineVersion || '',
    error: doc.error || null,
  };
}

// ============================================
// MAIN QUERY FUNCTION
// ============================================

export async function getAiGenerations(params: {
  companyId: string;
  moduleSource?: string;
  status?: string;
  provider?: string;
  startDate?: string;
  endDate?: string;
  page?: number;
  limit?: number;
  sort?: string;
  search?: string;
}): Promise<AiGenerationListResponse> {
  const { AiContext } = getModels();

  const {
    companyId,
    moduleSource,
    status,
    provider,
    startDate,
    endDate,
    page = 1,
    limit = 20,
    sort = 'createdAt',
    search,
  } = params;

  // Build query filter
  const filter: any = { companyId };

  if (moduleSource && moduleSource !== 'all') {
    filter.moduleSource = moduleSource;
  }

  if (provider && provider !== 'all') {
    filter.provider = provider;
  }

  if (status && status !== 'all') {
    // Map dashboard statuses back to AiContext statuses
    if (status === 'completed') {
      filter.status = { $in: ['draft', 'reviewed', 'approved'] };
    } else if (status === 'failed') {
      filter.status = 'rejected';
    } else {
      filter.status = status;
    }
  }

  // Date range filter. The UI sends date-only strings ("YYYY-MM-DD"), which
  // `new Date()` parses as UTC midnight. startDate is used as the start of the
  // selected day; endDate must be pushed to the END of the selected day
  // (23:59:59.999 UTC), otherwise `$lte: midnight` wrongly excludes every record
  // created after 00:00 on that day — which made the filter appear broken and
  // hid same-day results. A future startDate therefore yields an empty set
  // (no historical records leak in) rather than falling back to old data.
  if (startDate) {
    const start = new Date(startDate);
    if (!isNaN(start.getTime())) {
      filter.createdAt = { ...filter.createdAt, $gte: start };
    }
  }
  if (endDate) {
    const end = new Date(endDate);
    if (!isNaN(end.getTime())) {
      end.setUTCHours(23, 59, 59, 999);
      filter.createdAt = { ...filter.createdAt, $lte: end };
    }
  }

  // Trim + escape regex metacharacters so arbitrary input (e.g. "gpt-4o (v2)",
  // "a+b") is matched literally instead of being compiled as an invalid regex
  // (which otherwise throws a 500 and returns no/incorrect results).
  const searchTerm = typeof search === 'string' ? search.trim() : '';
  if (searchTerm) {
    const safe = escapeRegex(searchTerm);
    filter.$or = [
      { moduleSource: { $regex: safe, $options: 'i' } },
      { provider: { $regex: safe, $options: 'i' } },
      { aiModel: { $regex: safe, $options: 'i' } },
      { 'inputs.companyName': { $regex: safe, $options: 'i' } },
    ];
  }

  // Determine sort direction
  const sortDir = sort.startsWith('-') ? -1 : 1;
  const sortField = sort.replace(/^-/, '');
  const sortObj: any = {};
  sortObj[sortField] = sortDir;

  // Execute query with pagination
  const skip = (page - 1) * limit;
  const [docs, total] = await Promise.all([
    AiContext.find(filter).sort(sortObj).skip(skip).limit(limit).lean(),
    AiContext.countDocuments(filter),
  ]);

  // Compute aggregate stats (including token totals).
  //
  // Matched on the SAME `filter` as the list query above, not on `companyId`
  // alone. With companyId-only matching the cards ignored every module, status,
  // provider, date and search filter, so narrowing the table left the counts
  // unchanged and they disagreed with what was on screen.
  //
  // The completed/failed buckets mirror formatRecord's status mapping
  // (draft|reviewed|approved → completed, rejected → failed), so the counts
  // describe the same statuses the rows display.
  const statsAgg = await AiContext.aggregate([
    { $match: filter },
    {
      $group: {
        _id: null,
        totalGenerations: { $sum: 1 },
        completedCount: {
          $sum: {
            $cond: [{ $in: ['$status', ['draft', 'reviewed', 'approved']] }, 1, 0],
          },
        },
        failedCount: {
          $sum: { $cond: [{ $eq: ['$status', 'rejected'] }, 1, 0] },
        },
        avgProcessingTimeMs: { $avg: '$processingTimeMs' },
        // Summed for the dashboard's Total Duration card.
        totalProcessingTimeMs: { $sum: { $ifNull: ['$processingTimeMs', 0] } },
        totalInputTokens: { $sum: { $ifNull: ['$inputTokens', 0] } },
        totalOutputTokens: { $sum: { $ifNull: ['$outputTokens', 0] } },
        totalTokensUsed: { $sum: { $ifNull: ['$tokensUsed', 0] } },
      },
    },
  ]);

  const providerAgg = await AiContext.aggregate([
    { $match: { companyId } },
    { $group: { _id: '$provider', count: { $sum: 1 } } },
  ]);

  const moduleAgg = await AiContext.aggregate([
    { $match: { companyId } },
    { $group: { _id: '$moduleSource', count: { $sum: 1 } } },
  ]);

  const statsData = statsAgg[0] || {
    totalGenerations: 0,
    completedCount: 0,
    failedCount: 0,
    avgProcessingTimeMs: null,
    totalProcessingTimeMs: 0,
    totalInputTokens: 0,
    totalOutputTokens: 0,
    totalTokensUsed: 0,
  };

  // Estimate total cost across all records
  const estimatedCost = await estimateTotalCost(companyId);

  const providerBreakdown: Record<string, number> = {};
  for (const p of providerAgg) {
    if (p._id) providerBreakdown[p._id] = p.count;
  }

  const moduleBreakdown: Record<string, number> = {};
  for (const m of moduleAgg) {
    if (m._id) moduleBreakdown[m._id] = m.count;
  }

  const records = docs.map(formatRecord);

  return {
    records,
    total,
    page,
    limit,
    totalPages: Math.ceil(total / limit),
    stats: {
      totalGenerations: statsData.totalGenerations,
      completedCount: statsData.completedCount,
      failedCount: statsData.failedCount,
      processingCount: 0, // Active jobs come from aiJobStore, not DB
      avgProcessingTimeMs: statsData.avgProcessingTimeMs ?? null,
      totalProcessingTimeMs: statsData.totalProcessingTimeMs ?? 0,
      totalInputTokens: statsData.totalInputTokens,
      totalOutputTokens: statsData.totalOutputTokens,
      totalTokensUsed: statsData.totalTokensUsed,
      estimatedCost,
      providerBreakdown,
      moduleBreakdown,
    },
  };
}

// ============================================
// SINGLE RECORD DETAIL
// ============================================

export async function getAiGenerationDetail(id: string): Promise<AiGenerationDetailRecord | null> {
  const { AiContext } = getModels();
  const doc = await AiContext.findById(id).lean();
  if (!doc) return null;
  const record = formatRecord(doc);
  return {
    ...record,
    analysis: doc.analysis || null,
  };
}