import { eq } from 'drizzle-orm';
import { db } from '../db.js';
import { recommendations } from '../db/schema.js';
import { runInterpreter } from '../agents/interpreter.js';
import { runRetrievalBatch } from '../agents/retrievalContinuous.js';
import { runValidator } from '../agents/validator.js';
import { runRanking } from '../agents/ranking.js';
import { runCurator } from '../agents/curator.js';
import type { CuratorOutput, InterpreterOutput, MediaType, RankingOutput, RetrievalCandidate, SSEEvent } from '../types/agents.js';
import { generateTitle } from '../agents/titleGenerator.js';

function log(msg: string, data?: unknown) {
  const ts = new Date().toISOString();
  if (data !== undefined) {
    console.log(`[continuous-pipeline] [${ts}] ${msg}`, data);
  } else {
    console.log(`[continuous-pipeline] [${ts}] ${msg}`);
  }
}

function deduplicateCandidates(candidates: RetrievalCandidate[], seenTitles: Set<string>): RetrievalCandidate[] {
  return candidates.filter((c) => {
    const key = c.title.toLowerCase();
    if (seenTitles.has(key)) return false;
    seenTitles.add(key);
    return true;
  });
}

function splitIntoBuckets<T>(items: T[], n: number): T[][] {
  const size = Math.ceil(items.length / n);
  return Array.from({ length: n }, (_, i) => items.slice(i * size, (i + 1) * size))
    .filter((b) => b.length > 0);
}

function mergeCuratorOutputs(a: CuratorOutput[], b: CuratorOutput[]): CuratorOutput[] {
  const seen = new Set(a.map((x) => x.title.toLowerCase()));
  return [...a, ...b.filter((x) => !seen.has(x.title.toLowerCase()))];
}

interface ContinuousPipelineInput {
  likedShows: string;
  dislikedShows?: string;
  themes?: string;
  requirements?: string;
  avoid?: string;
  totalCount: number;
  useWebSearch: boolean;
  validateResults: boolean;
  mediaType: MediaType;
}

export async function runContinuousPipeline(
  recId: string,
  input: ContinuousPipelineInput,
  sseWrite: (event: SSEEvent) => void,
): Promise<CuratorOutput[]> {
  const startTime = Date.now();
  const {
    likedShows,
    dislikedShows = '',
    themes = '',
    requirements = '',
    avoid = '',
    totalCount,
    useWebSearch,
    validateResults,
    mediaType,
  } = input;

  const totalBatches = Math.ceil(totalCount / 10);
  const allSeenTitles = new Set<string>();
  let previousResponseId: string | null = null;
  let interpreterOutput: InterpreterOutput | null = null;
  const accumulatedCandidates: RetrievalCandidate[] = [];

  log(`Starting continuous pipeline: ${totalCount} titles (${totalBatches} batches), mediaType=${mediaType}, webSearch=${useWebSearch}, validate=${validateResults}`);

  try {
    // --- Interpreter (runs once) ---
    log('Interpreter: start');
    sseWrite({ stage: 'interpreter', status: 'start' });
    const t0 = Date.now();

    interpreterOutput = await runInterpreter({
      main_prompt: themes || 'recommend shows based on user preferences',
      liked_shows: likedShows,
      disliked_shows: dislikedShows,
      themes: themes,
      media_type: mediaType,
    });

    log(`Interpreter: done (${Date.now() - t0}ms)`, {
      liked: interpreterOutput.liked,
      themes: interpreterOutput.themes,
    });
    sseWrite({ stage: 'interpreter', status: 'done', data: interpreterOutput });

    // --- Retrieval (batched, chained) ---
    log(`Retrieval: start (${totalBatches} batches)`);
    sseWrite({ stage: 'retrieval', status: 'start' });

    const hardRequirements = requirements.length > 0;

    for (let batchIndex = 0; batchIndex < totalBatches; batchIndex++) {
      log(`Retrieval batch ${batchIndex + 1}/${totalBatches}: start`);
      sseWrite({
        stage: 'retrieval',
        status: 'start',
        data: {
          batch: batchIndex + 1,
          totalBatches,
          totalCandidates: accumulatedCandidates.length,
        },
      });

      const seenTitlesArray = Array.from(allSeenTitles);
      const result = await runRetrievalBatch(
        interpreterOutput,
        mediaType,
        useWebSearch,
        hardRequirements,
        seenTitlesArray,
        previousResponseId
      );

      previousResponseId = result.responseId;

      // Deduplicate against previously seen titles
      const deduped = deduplicateCandidates(result.candidates, allSeenTitles);
      accumulatedCandidates.push(...deduped);

      log(`Retrieval batch ${batchIndex + 1}/${totalBatches}: done, ${deduped.length} new candidates (total: ${accumulatedCandidates.length})`);
    }

    log(`Retrieval: complete, ${accumulatedCandidates.length} total candidates`);
    sseWrite({
      stage: 'retrieval',
      status: 'done',
      data: {
        totalBatches,
        totalCandidates: accumulatedCandidates.length,
      },
    });

    // --- Validator (optional, per batch already done during retrieval) ---
    // Note: In continuous mode, we could run validator after each batch
    // For now, we'll run it once after all batches if validateResults is true
    let candidatesForRanking = accumulatedCandidates;

    if (validateResults && candidatesForRanking.length > 0) {
      log(`Validator: start (${candidatesForRanking.length} candidates)`);
      sseWrite({ stage: 'validator', status: 'start' });

      const tV = Date.now();
      const validatorOutput = await runValidator(candidatesForRanking, mediaType);
      const verified = validatorOutput.candidates.filter((c) => !c.isTrash);
      const trashCount = validatorOutput.candidates.length - verified.length;

      candidatesForRanking = verified.map(({ title, reason }) => ({ title, reason }));
      allSeenTitles.clear();
      candidatesForRanking.forEach((c) => allSeenTitles.add(c.title.toLowerCase()));

      log(`Validator: done (${Date.now() - tV}ms) — removed ${trashCount} trash entries`);
      sseWrite({ stage: 'validator', status: 'done', data: { removed: trashCount, remaining: candidatesForRanking.length } });
    } else {
      sseWrite({ stage: 'validator', status: 'done', data: { skipped: !validateResults } });
    }

    // --- Ranking (bucketed) ---
    log(`Ranking: start (${candidatesForRanking.length} candidates)`);
    sseWrite({ stage: 'ranking', status: 'start' });
    const t2 = Date.now();

    const rankBucketCount = Math.ceil(candidatesForRanking.length / 15) || 1;
    const candidateBuckets = splitIntoBuckets(candidatesForRanking, rankBucketCount);

    const rankingBuckets = await Promise.all(
      candidateBuckets.map((bucket) =>
        runRanking(interpreterOutput!, { candidates: bucket }, mediaType, hardRequirements)
      )
    );

    const dedupTitles = (titles: string[]) => [...new Map(titles.map((t) => [t.toLowerCase(), t])).values()];
    const rankingOutput: RankingOutput = {
      full_match: dedupTitles(rankingBuckets.flatMap((r) => r.full_match)),
      definitely_like: dedupTitles(rankingBuckets.flatMap((r) => r.definitely_like)),
      might_like: dedupTitles(rankingBuckets.flatMap((r) => r.might_like)),
      questionable: dedupTitles(rankingBuckets.flatMap((r) => r.questionable)),
      will_not_like: dedupTitles(rankingBuckets.flatMap((r) => r.will_not_like)),
    };

    log(`Ranking: done (${Date.now() - t2}ms)`, {
      full_match: rankingOutput.full_match.length,
      definitely_like: rankingOutput.definitely_like.length,
      might_like: rankingOutput.might_like.length,
    });
    sseWrite({ stage: 'ranking', status: 'done', data: rankingOutput });

    // print all ranked titles for debuggings
    log('Ranked titles:', {
      full_match: rankingOutput.full_match,
      definitely_like: rankingOutput.definitely_like,
      might_like: rankingOutput.might_like,
      questionable: rankingOutput.questionable,
      will_not_like: rankingOutput.will_not_like,
    });

    // --- Curator (bucketed) ---
    log(`Curator: start`);
    sseWrite({ stage: 'curator', status: 'start' });
    const t3 = Date.now();

    type CategorizedItem = { title: string; category: keyof RankingOutput };
    const categorizedItems: CategorizedItem[] = [
      ...rankingOutput.full_match.map((t) => ({ title: t, category: 'full_match' as const })),
      ...rankingOutput.definitely_like.map((t) => ({ title: t, category: 'definitely_like' as const })),
      ...rankingOutput.might_like.map((t) => ({ title: t, category: 'might_like' as const })),
      ...rankingOutput.questionable.map((t) => ({ title: t, category: 'questionable' as const })),
      ...rankingOutput.will_not_like.map((t) => ({ title: t, category: 'will_not_like' as const })),
    ];

    const curatorBucketCount = Math.ceil(categorizedItems.length / 15) || 1;
    const curatorItemBuckets = splitIntoBuckets(categorizedItems, curatorBucketCount);
    const curatorBucketRankings: RankingOutput[] = curatorItemBuckets.map((bucket) => ({
      full_match: bucket.filter((i) => i.category === 'full_match').map((i) => i.title),
      definitely_like: bucket.filter((i) => i.category === 'definitely_like').map((i) => i.title),
      might_like: bucket.filter((i) => i.category === 'might_like').map((i) => i.title),
      questionable: bucket.filter((i) => i.category === 'questionable').map((i) => i.title),
      will_not_like: bucket.filter((i) => i.category === 'will_not_like').map((i) => i.title),
    }));

    const curatorBucketOutputs = await Promise.all(
      curatorBucketRankings.map((ranking) =>
        runCurator(ranking, interpreterOutput!, mediaType, useWebSearch)
      )
    );

    const curatorOutput = curatorBucketOutputs.reduce((acc, bucket) => mergeCuratorOutputs(acc, bucket), [] as CuratorOutput[]);
    log(`Curator: done (${Date.now() - t3}ms) — ${curatorOutput.length} items curated`);
    sseWrite({ stage: 'curator', status: 'done', data: curatorOutput });

    // Generate AI title
    let aiTitle = 'Continuous Recommendations';
    try {
      log('Title generation: start');
      aiTitle = await generateTitle(interpreterOutput, mediaType);
      log(`Title generation: done — "${aiTitle}"`);
    } catch (err) {
      log(`Title generation failed: ${String(err)}`);
    }

    // Sort by category order
    const CATEGORY_ORDER: Record<string, number> = {
      'Full Match': 0,
      'Definitely Like': 1,
      'Might Like': 2,
      'Questionable': 3,
      'Will Not Like': 4,
    };
    curatorOutput.sort((a, b) => (CATEGORY_ORDER[a.category] ?? 99) - (CATEGORY_ORDER[b.category] ?? 99));

    // Save to database (update existing record)
    log('Saving results to DB');
    await db
      .update(recommendations)
      .set({
        title: aiTitle,
        main_prompt: themes || 'Continuous recommendations',
        liked_shows: likedShows,
        disliked_shows: dislikedShows,
        themes: themes,
        brainstorm_count: totalCount,
        media_type: mediaType,
        use_web_search: useWebSearch,
        use_validator: validateResults,
        status: 'done',
        recommendations: curatorOutput,
      })
      .where(eq(recommendations.id, recId));

    sseWrite({ stage: 'complete', status: 'done', data: { id: recId, title: aiTitle, batchCount: totalBatches } });

    log(`Pipeline complete (total: ${Date.now() - startTime}ms), saved as ${recId}`);
    return curatorOutput;

  } catch (err) {
    const message = err instanceof Error ? err.message : String(err);
    log(`Pipeline error: ${message}`);

    // Update status to error
    await db
      .update(recommendations)
      .set({ status: 'error' })
      .where(eq(recommendations.id, recId));

    sseWrite({ stage: 'curator', status: 'error', data: { message } });
    return [];
  }
}
