import type {
  Task,
  TaskStage,
  TaskStageEvent,
  TaskType
} from '@/lib/types';

export type StageTimelineItem = {
  stage: TaskStage;
  label: string;
  state: 'completed' | 'active' | 'pending';
  detail: string | null;
  timestamp: string | null;
  context: Task['stage_context'] | null;
};

const TASK_TYPE_LABELS: Record<TaskType, string> = {
  sync_filings: 'Filing sync',
  refresh_prices: 'Price refresh',
  analyze_filing: 'Filing analysis',
  portfolio_insights: 'Portfolio insight',
  index_search: 'Search indexing'
};

const STAGE_LABELS: Record<TaskStage, string> = {
  queued: 'Queued',
  running: 'Running',
  completed: 'Completed',
  failed: 'Failed',
  'sync.fetch_filings': 'Fetch filings',
  'sync.discover_assets': 'Discover taxonomy assets',
  'sync.extract_taxonomy': 'Extract taxonomy',
  'sync.normalize_taxonomy': 'Normalize taxonomy',
  'sync.derive_metrics': 'Derive metrics',
  'sync.validate_pdf_metrics': 'Validate PDF metrics',
  'sync.persist_taxonomy': 'Persist taxonomy',
  'sync.fetch_metrics': 'Fetch filing metrics',
  'sync.persist_filings': 'Persist filings',
  'sync.hydrate_statements': 'Hydrate statements',
  'refresh.load_holdings': 'Load holdings',
  'refresh.fetch_quotes': 'Fetch quotes',
  'refresh.persist_prices': 'Persist prices',
  'analyze.load_filing': 'Load filing',
  'analyze.fetch_document': 'Fetch primary document',
  'analyze.extract': 'Extract context',
  'analyze.generate_report': 'Generate report',
  'analyze.persist_report': 'Persist report',
  'search.collect_sources': 'Collect sources',
  'search.fetch_documents': 'Fetch documents',
  'search.chunk': 'Chunk content',
  'search.embed': 'Generate embeddings',
  'search.persist': 'Persist search index',
  'insights.load_holdings': 'Load holdings',
  'insights.generate': 'Generate insight',
  'insights.persist': 'Persist insight'
};

const TASK_STAGE_ORDER: Record<TaskType, TaskStage[]> = {
  sync_filings: [
    'queued',
    'running',
    'sync.fetch_filings',
    'sync.persist_filings',
    'sync.discover_assets',
    'sync.extract_taxonomy',
    'sync.normalize_taxonomy',
    'sync.derive_metrics',
    'sync.validate_pdf_metrics',
    'sync.persist_taxonomy',
    'completed'
  ],
  refresh_prices: [
    'queued',
    'running',
    'refresh.load_holdings',
    'refresh.fetch_quotes',
    'refresh.persist_prices',
    'completed'
  ],
  analyze_filing: [
    'queued',
    'running',
    'analyze.load_filing',
    'analyze.fetch_document',
    'analyze.extract',
    'analyze.generate_report',
    'analyze.persist_report',
    'completed'
  ],
  index_search: [
    'queued',
    'running',
    'search.collect_sources',
    'search.fetch_documents',
    'search.chunk',
    'search.embed',
    'search.persist',
    'completed'
  ],
  portfolio_insights: [
    'queued',
    'running',
    'insights.load_holdings',
    'insights.generate',
    'insights.persist',
    'completed'
  ]
};

export function taskTypeLabel(taskType: TaskType) {
  return TASK_TYPE_LABELS[taskType];
}

export function stageLabel(stage: TaskStage) {
  return STAGE_LABELS[stage] ?? stage;
}

export function taskStageOrder(taskType: TaskType) {
  return TASK_STAGE_ORDER[taskType] ?? ['queued', 'running', 'completed'];
}

export function fallbackStageProgress(task: Pick<Task, 'task_type' | 'stage' | 'status'>) {
  const orderedStages = taskStageOrder(task.task_type);
  const stageIndex = orderedStages.indexOf(task.stage);

  if (stageIndex === -1) {
    return null;
  }

  if (task.status === 'completed') {
    return {
      current: orderedStages.length,
      total: orderedStages.length,
      unit: 'steps'
    };
  }

  if (task.status === 'failed') {
    return {
      current: Math.max(stageIndex + 1, 1),
      total: orderedStages.length,
      unit: 'steps'
    };
  }

  return {
    current: Math.max(stageIndex + 1, 1),
    total: orderedStages.length,
    unit: 'steps'
  };
}

export function buildStageTimeline(task: Task, events: TaskStageEvent[]): StageTimelineItem[] {
  const baseOrder = taskStageOrder(task.task_type);
  const orderedStages = [...baseOrder];

  if (task.status === 'failed' && !orderedStages.includes('failed')) {
    orderedStages.push('failed');
  }

  const latestEventByStage = new Map<TaskStage, TaskStageEvent>();
  for (const event of events) {
    latestEventByStage.set(event.stage, event);
  }

  return orderedStages.map((stage) => {
    const event = latestEventByStage.get(stage);

    if (task.status === 'queued' || task.status === 'running') {
      if (stage === task.stage) {
        return {
          stage,
          label: stageLabel(stage),
          state: 'active' as const,
          detail: event?.stage_detail ?? task.stage_detail,
          timestamp: event?.created_at ?? task.updated_at,
          context: task.stage_context ?? event?.stage_context ?? null
        };
      }

      if (event) {
        return {
          stage,
          label: stageLabel(stage),
          state: 'completed' as const,
          detail: event.stage_detail,
          timestamp: event.created_at,
          context: event.stage_context ?? null
        };
      }

      return {
        stage,
        label: stageLabel(stage),
        state: 'pending' as const,
        detail: null,
        timestamp: null,
        context: null
      };
    }

    if (stage === task.stage || event) {
      return {
        stage,
        label: stageLabel(stage),
        state: 'completed' as const,
        detail: event?.stage_detail ?? task.stage_detail,
        timestamp: event?.created_at ?? task.finished_at,
        context: event?.stage_context ?? (stage === task.stage ? task.stage_context : null) ?? null
      };
    }

    return {
      stage,
      label: stageLabel(stage),
      state: 'pending' as const,
      detail: null,
      timestamp: null,
      context: null
    };
  });
}