import { AGENT_TYPES, AgentEvent, AgentPayload, TraceSource } from '../lib/types/agent';
import { ReasoningProfile } from '../lib/types/llm';
import { logger } from '../lib/logger';
import { Context } from 'aws-lambda';
import { validateEventPayload, buildProcessOptions, initAgent } from '../lib/utils/agent-helpers';
import { emitTaskEvent } from '../lib/utils/agent-helpers/event-emitter';
import { SWARM } from '../lib/constants';
import { RESEARCH_TASK_METADATA } from '../lib/schema/events';

/**
 * Technical Research Agent Handler.
 * Executes deep technical research by discovering patterns, reading docs,
 * and analyzing codebases using specialized MCP tools.
 *
 * Implements the 5-phase operational protocol: Discovery, Mapping, Analysis, Synthesis, Recommendation.
 *
 * @param event - The EventBridge event.
 * @param context - The AWS Lambda context.
 * @returns A promise that resolves to the agent's response string, or undefined on error.
 */
export const handler = async (event: AgentEvent, context: Context): Promise<string | undefined> => {
  logger.info('Researcher Agent received task:', JSON.stringify(event, null, 2));

  const payload = validateEventPayload<AgentPayload>(event, `${AGENT_TYPES.RESEARCHER}_task`);
  const {
    userId,
    task,
    metadata,
    traceId,
    sessionId,
    isContinuation,
    initiatorId,
    depth,
    tokenBudget,
    costLimit,
    workspaceId,
  } = payload;

  const isAggregation = task?.includes('[AGGREGATED_RESULTS]');

  // 1. Discovery & Initialization (config + context loaded in parallel)
  const { memory, agent } = await initAgent(AGENT_TYPES.RESEARCHER, { workspaceId });

  // 2. Swarm Self-Organization: Decompose high-level goals into parallel exploration
  if (!isAggregation && (depth ?? 0) < SWARM.MAX_RECURSIVE_DEPTH && task) {
    const { decomposePlan } = await import('../lib/agent/decomposer');
    const parentGapIds =
      (metadata?.gapIds as string[]) ?? (metadata?.coveredGapIds as string[]) ?? [];
    const decomposed = await decomposePlan(
      task,
      payload.taskId || traceId || `plan-${Date.now()}`,
      parentGapIds,
      {
        defaultAgentId: AGENT_TYPES.RESEARCHER,
        maxSubTasks: SWARM.DEFAULT_MAX_SUB_TASKS,
        minLength: 300,
      }
    );

    if (decomposed.wasDecomposed && decomposed.subTasks.length > 1) {
      logger.info(
        `[RESEARCHER] Goal detected. Decomposing into ${decomposed.subTasks.length} parallel sub-tasks.`
      );

      const { emitTypedEvent } = await import('../lib/utils/typed-emit');
      const { EventType } = await import('../lib/types/agent');

      const subTaskEvents = decomposed.subTasks.map((sub) => ({
        taskId: sub.subTaskId,
        agentId: sub.agentId,
        task: sub.task,
        metadata: {
          ...metadata,
          traceId: traceId ?? sub.planId,
          subTaskId: sub.subTaskId,
          planId: sub.planId,
        },
        tokenBudget: tokenBudget ? Math.floor(tokenBudget / decomposed.subTasks.length) : undefined,
        costLimit: costLimit ? costLimit / decomposed.subTasks.length : undefined,
      }));

      try {
        await emitTypedEvent(AGENT_TYPES.RESEARCHER, EventType.PARALLEL_TASK_DISPATCH, {
          userId: userId,
          tasks: subTaskEvents,
          barrierTimeoutMs: 10 * 60 * 1000, // 10 mins
          aggregationType: 'agent_guided',
          aggregationPrompt: `I have completed the parallel exploration for: "${task}". 
                             Synthesize the findings from these sub-tasks into a cohesive technical report. 
                             Focus on pattern discovery and actionable technical debt/gap identification.
                             Prepend the response with [AGGREGATED_RESULTS].`,
          initialQuery: task,
          traceId,
          initiatorId: AGENT_TYPES.RESEARCHER,
          depth: depth ?? 0,
          sessionId,
          tokenBudget,
          costLimit,
        });

        // Return a PAUSED signal
        return `TASK_PAUSED: Decomposed research goal into ${decomposed.subTasks.length} parallel sub-tasks.`;
      } catch (dispatchError) {
        logger.error(`[RESEARCHER] Failed to dispatch parallel tasks:`, dispatchError);
        return `TASK_FAILED: Parallel dispatch failed. Falling back to local execution.`;
      }
    }
  }

  // 3. Execution
  const parsedMetadata = RESEARCH_TASK_METADATA.parse(metadata || {});
  const timeBudgetMs = parsedMetadata.timeBudgetMs ?? 600000;

  const gapIds = (metadata?.gapIds as string[]) ?? (metadata?.coveredGapIds as string[]) ?? [];
  const specDetails: string[] = [];
  if (gapIds && gapIds.length > 0) {
    for (const gapId of gapIds) {
      try {
        const planStr = await memory.getDistilledMemory(`PLAN#${gapId}`);
        if (planStr) {
          const parsedPlan = JSON.parse(planStr);
          if (parsedPlan.spec) {
            specDetails.push(`### Target Spec for Gap [${gapId}]:\n${parsedPlan.spec}`);
          }
        }
      } catch (e) {
        logger.warn(`[RESEARCHER] Failed to load spec for gap ${gapId}:`, e);
      }
    }
  }

  const specContext =
    specDetails.length > 0
      ? `\n\n[TARGET_TECHNICAL_SPECIFICATIONS]:\nThese are the active technical specifications and EARS evolution contracts for this research scope. Use them to focus your mapping and discovery:\n${specDetails.join('\n\n')}`
      : '';

  const researchTask = `${task || ''}${specContext}`;

  const processOptions = buildProcessOptions({
    isContinuation,
    isIsolated: true,
    initiatorId,
    depth,
    traceId,
    taskId: payload.taskId,
    sessionId,
    source: (payload.source as TraceSource) || TraceSource.UNKNOWN,
    context,
    taskTimeoutMs: timeBudgetMs,
    profile: ReasoningProfile.THINKING,
    tokenBudget,
    costLimit,
  });

  const startTime = Date.now();
  let finalResponseText: string;

  try {
    const result = await agent.process(userId, researchTask, processOptions);
    finalResponseText = result?.responseText || '';

    // 4. Persistence: Store final synthesis as a granular memory item
    const existingFindings = await memory.searchInsights(
      userId,
      finalResponseText.substring(0, 100),
      undefined,
      1
    );

    const isDuplicate = existingFindings?.items.some(
      (item) => item.tags?.includes(traceId ?? '') || item.tags?.includes(payload.taskId ?? '')
    );

    if (!isDuplicate) {
      await memory.addMemory(userId, 'research_finding', finalResponseText, {
        category: 'research_finding',
        confidence: 9,
        impact: 5,
        tags: [traceId, payload.taskId, 'synthesis'].filter((t): t is string => !!t),
      });
    }

    // 5. Notification
    await emitTaskEvent({
      source: `${AGENT_TYPES.RESEARCHER}.agent`,
      agentId: AGENT_TYPES.RESEARCHER,
      userId: userId,
      task: task || '',
      response: finalResponseText,
      traceId,
      taskId: payload.taskId,
      sessionId,
      initiatorId,
      depth,
      metadata: {
        durationMs: Date.now() - startTime,
        findingsCategory: 'research_finding',
      },
    });

    return finalResponseText;
  } catch (error) {
    logger.error(`[RESEARCHER] Task failed:`, error);
    const errorMsg = error instanceof Error ? error.message : String(error);

    await emitTaskEvent({
      source: `${AGENT_TYPES.RESEARCHER}.agent`,
      agentId: AGENT_TYPES.RESEARCHER,
      userId: userId,
      task: task || '',
      error: `Research task failed: ${errorMsg}`,
      traceId,
      taskId: payload.taskId,
      sessionId,
      initiatorId,
      depth,
    });

    throw error;
  }
};
