import { ResearcherOutput, SearchAgentInput } from './types';
import SessionManager from '@/lib/session';
import { classify } from './classifier';
import Researcher from './researcher';
import { getWriterPrompt } from '@/lib/prompts/search/writer';
import { WidgetExecutor } from './widgets';
import db from '@/lib/db';
import { chats, messages } from '@/lib/db/schema';
import { and, eq, gt } from 'drizzle-orm';
import { TextBlock } from '@/lib/types';

class SearchAgent {
  async searchAsync(session: SessionManager, input: SearchAgentInput) {
    try {
      await this.doSearch(session, input);
    } catch (err) {
      console.error('[SearchAgent] Fatal:', err);
      const blocks = session.getAllBlocks();
      const sourceBlock = blocks.find((b) => b.type === 'source');
      let sources: { metadata?: { title?: string } }[] =
        sourceBlock?.type === 'source' ? sourceBlock.data : [];
      if (sources.length === 0) {
        const researchBlock = blocks.find(
          (b): b is typeof b & { data: { subSteps?: { type: string; reading?: { metadata?: { title?: string } }[] }[] } } =>
            b.type === 'research' && 'subSteps' in (b.data ?? {}),
        );
        const searchStep = researchBlock?.data?.subSteps?.find(
          (s) => s.type === 'search_results' && Array.isArray(s.reading),
        );
        if (searchStep && 'reading' in searchStep) {
          sources = searchStep.reading ?? [];
        }
      }
      if (sources.length > 0) {
        const lines = sources.slice(0, 10).map(
          (s: { metadata?: { title?: string } }, i: number) =>
            `${i + 1}. **${s?.metadata?.title ?? 'Источник'}**`,
        );
        session.emitBlock({
          id: crypto.randomUUID(),
          type: 'text',
          data: `## По найденным источникам\n\n${lines.join('\n')}\n\n*Ответ LLM недоступен (400). Проверьте модель gemini-3-flash в Settings или попробуйте другую модель.*`,
        });
      } else {
        session.emitBlock({
          id: crypto.randomUUID(),
          type: 'text',
          data: `Ошибка: ${err instanceof Error ? err.message : String(err)}. Проверьте LLM в Settings.`,
        });
      }
      session.emit('end', {});
    }
  }

  private async doSearch(session: SessionManager, input: SearchAgentInput) {
    const exists = await db.query.messages.findFirst({
      where: and(
        eq(messages.chatId, input.chatId),
        eq(messages.messageId, input.messageId),
      ),
    });

    if (!exists) {
      await db.insert(messages).values({
        chatId: input.chatId,
        messageId: input.messageId,
        backendId: session.id,
        query: input.followUp,
        createdAt: new Date().toISOString(),
        status: 'answering',
        responseBlocks: [],
      });
    } else {
      await db
        .delete(messages)
        .where(
          and(eq(messages.chatId, input.chatId), gt(messages.id, exists.id)),
        )
        .execute();
      await db
        .update(messages)
        .set({
          status: 'answering',
          backendId: session.id,
          responseBlocks: [],
        })
        .where(
          and(
            eq(messages.chatId, input.chatId),
            eq(messages.messageId, input.messageId),
          ),
        )
        .execute();
    }

    const classification = await classify({
      chatHistory: input.chatHistory,
      enabledSources: input.config.sources,
      query: input.followUp,
      llm: input.config.llm,
      locale: input.config.locale,
    });

    const widgetPromise = WidgetExecutor.executeAll({
      classification,
      chatHistory: input.chatHistory,
      followUp: input.followUp,
      llm: input.config.llm,
    }).then((widgetOutputs) => {
      widgetOutputs.forEach((o) => {
        session.emitBlock({
          id: crypto.randomUUID(),
          type: 'widget',
          data: {
            widgetType: o.type,
            params: o.data,
          },
        });
      });
      return widgetOutputs;
    });

    let searchPromise: Promise<ResearcherOutput> | null = null;

    if (!classification.classification.skipSearch) {
      const researcher = new Researcher();
      searchPromise = researcher.research(session, {
        chatHistory: input.chatHistory,
        followUp: input.followUp,
        classification: classification,
        config: input.config,
      });
    }

    const [widgetOutputs, searchResults] = await Promise.all([
      widgetPromise,
      searchPromise,
    ]);

    session.emit('data', {
      type: 'researchComplete',
    });

    const MAX_RESULTS_FOR_WRITER = 15;
    const MAX_CONTENT_PER_RESULT = 180;
    const findingsForWriter =
      searchResults?.searchFindings.slice(0, MAX_RESULTS_FOR_WRITER) ?? [];
    const finalContext =
      findingsForWriter
        .map((f, index) => {
          const content =
            f.content.length > MAX_CONTENT_PER_RESULT
              ? f.content.slice(0, MAX_CONTENT_PER_RESULT) + '…'
              : f.content;
          return `<result index=${index + 1} title="${String(f.metadata.title).replace(/"/g, "'")}">${content}</result>`;
        })
        .join('\n') || '';

    const widgetContext = widgetOutputs
      .map((o) => {
        return `<result>${o.llmContext}</result>`;
      })
      .join('\n-------------\n');

    const finalContextWithWidgets = `<search_results note="These are the search results and assistant can cite these">\n${finalContext}\n</search_results>\n<widgets_result noteForAssistant="Its output is already showed to the user, assistant can use this information to answer the query but do not CITE this as a souce">\n${widgetContext}\n</widgets_result>`;

    const writerPrompt = getWriterPrompt(
      finalContextWithWidgets,
      input.config.systemInstructions,
      input.config.mode,
      input.config.locale,
    );

    const answerStream = input.config.llm.streamText({
      messages: [
        { role: 'system', content: writerPrompt },
        ...input.chatHistory,
        { role: 'user', content: input.followUp },
      ],
      options: { maxTokens: 4096 },
    });

    let responseBlockId = '';
    let hasContent = false;

    for await (const chunk of answerStream) {
      if (!chunk.contentChunk && !responseBlockId) continue;
      if (!responseBlockId) {
        const block: TextBlock = {
          id: crypto.randomUUID(),
          type: 'text',
          data: chunk.contentChunk,
        };
        session.emitBlock(block);
        responseBlockId = block.id;
        if (chunk.contentChunk) hasContent = true;
      } else {
        const block = session.getBlock(responseBlockId) as TextBlock | null;
        if (block) {
          block.data += chunk.contentChunk;
          if (chunk.contentChunk) hasContent = true;
          session.updateBlock(block.id, [
            { op: 'replace', path: '/data', value: block.data },
          ]);
        }
      }
    }

    if (!hasContent && findingsForWriter.length > 0) {
      const lines = findingsForWriter.slice(0, 10).map((f, i) => {
        const title = f.metadata.title ?? 'Без названия';
        const excerpt =
          f.content.length > 120 ? f.content.slice(0, 120) + '…' : f.content;
        return `${i + 1}. **${title}** — ${excerpt}`;
      });
      session.emitBlock({
        id: crypto.randomUUID(),
        type: 'text',
        data: `## По найденным источникам\n\n${lines.join('\n\n')}\n\n*Ответ LLM недоступен. Проверьте модель в Settings.*`,
      });
    }

    session.emit('end', {});

    try {
      await db
        .update(messages)
        .set({
          status: 'completed',
          responseBlocks: session.getAllBlocks(),
        })
        .where(
          and(
            eq(messages.chatId, input.chatId),
            eq(messages.messageId, input.messageId),
          ),
        )
        .execute();
    } catch (dbErr) {
      console.error('[SearchAgent] DB update failed:', dbErr);
    }
  }
}

export default SearchAgent;
