import { HumanMessage, BaseMessage } from '@langchain/core/messages';
import type { TokenCounter } from '@/types/run';

export type SummarizeTier = 'full' | 'simple' | 'emergency';

export type SummarizeConfig = {
  /** Max tokens for the summary output */
  maxOutputTokens?: number;
  /** Token counter for validation */
  tokenCounter?: TokenCounter;
  /** Max tokens the summary should fit within */
  summaryBudget?: number;
  /** Whether this is a multi-agent conversation */
  isMultiAgent?: boolean;
  /** Current agent workflow state for multi-agent context */
  agentWorkflowState?: {
    currentAgentId: string;
    agentChain: string[];
    pendingAgents: string[];
  };
};

export type SummarizeResult = {
  summary: string;
  tier: SummarizeTier;
  tokenCount?: number;
  messagesCompacted: number;
};

const MAX_MESSAGE_CHARS = 4000;
const MAX_TRUNCATED_CHARS = 500;
const SIMPLE_MAX_TOKENS = 512;

export const FULL_SUMMARY_TEMPLATE = `Analyze this conversation and produce a structured summary covering each section below. Be precise and preserve actionable details.

## 1. Task Overview & User Intent
What the user originally asked for and their high-level goal.

## 2. Technical Context
Languages, frameworks, libraries, APIs, and architectural patterns involved.

## 3. Files & Code State
All file paths mentioned, their current state, and any pending modifications.

## 4. Problem Resolution History
Problems encountered, debugging steps taken, and their outcomes.

## 5. Agent Workflow State
{agent_workflow}

## 6. Tool Results Summary
Tools called, their inputs, and key outputs (omit verbose raw data).

## 7. Progress Tracking
What has been completed, what remains, and any blockers.

## 8. Active Working State
Variables, configurations, and runtime state that must be preserved.

## 9. Continuation Plan
Immediate next steps and the order in which they should be executed.

## 10. Critical Context
Any constraints, warnings, edge cases, or decisions that must not be lost.

Conversation:
{conversation}

Structured Summary:`;

export const SIMPLE_SUMMARY_TEMPLATE =
  "Summarize this conversation preserving: user's original request, key decisions, current task state, and any file paths mentioned.\n\nConversation:\n{conversation}\n\nSummary:";

export function formatMessagesForSummary(messages: BaseMessage[]): string {
  const parts: string[] = [];
  for (const message of messages) {
    const role = message.getType();
    const raw =
      typeof message.content === 'string'
        ? message.content
        : JSON.stringify(message.content);
    const content =
      raw.length > MAX_MESSAGE_CHARS
        ? raw.slice(0, MAX_MESSAGE_CHARS) + '...'
        : raw;
    parts.push(`${role}: ${content}`);
  }
  return parts.join('\n');
}

export function buildFullSummaryPrompt(
  conversation: string,
  config?: SummarizeConfig
): string {
  let agentWorkflow = 'N/A (single-agent conversation)';
  if (config?.isMultiAgent === true && config.agentWorkflowState != null) {
    const { currentAgentId, agentChain, pendingAgents } =
      config.agentWorkflowState;
    agentWorkflow = [
      `Current agent: ${currentAgentId}`,
      `Agent chain: ${agentChain.join(' -> ')}`,
      `Pending agents: ${pendingAgents.length > 0 ? pendingAgents.join(', ') : 'none'}`,
    ].join('\n');
  }

  return FULL_SUMMARY_TEMPLATE.replace(
    '{agent_workflow}',
    agentWorkflow
  ).replace('{conversation}', conversation);
}

export function buildSimpleSummaryPrompt(conversation: string): string {
  return SIMPLE_SUMMARY_TEMPLATE.replace('{conversation}', conversation);
}

export function createEmergencySummary(messages: BaseMessage[]): string {
  let firstUserContent = '';
  let lastAIContent = '';
  const toolNames = new Set<string>();

  for (const message of messages) {
    const type = message.getType();
    const content =
      typeof message.content === 'string'
        ? message.content
        : JSON.stringify(message.content);

    if (type === 'human' && firstUserContent.length === 0) {
      firstUserContent = content.slice(0, MAX_TRUNCATED_CHARS);
    }
    if (type === 'ai') {
      lastAIContent = content.slice(0, MAX_TRUNCATED_CHARS);
    }
    if (type === 'tool') {
      const name = message.name;
      if (name != null && name.length > 0) {
        toolNames.add(name);
      }
    }
  }

  const toolList = toolNames.size > 0 ? [...toolNames].join(', ') : 'none';
  return [
    '[Emergency Context Summary]',
    `Original request: ${firstUserContent}`,
    `Last response: ${lastAIContent}`,
    `Tools used: ${toolList}`,
    `Messages compacted: ${messages.length}`,
  ].join('\n');
}

export function validateSummarySize(
  summary: string,
  budget: number,
  tokenCounter: TokenCounter
): boolean {
  const count = tokenCounter(new HumanMessage(summary));
  return count <= budget;
}

export async function summarize(
  messages: BaseMessage[],
  callback: (prompt: string, maxTokens: number) => Promise<string | undefined>,
  config?: SummarizeConfig
): Promise<SummarizeResult> {
  const messageCount = messages.length;
  const maxOutputTokens = config?.maxOutputTokens ?? 1024;

  // Tier 1: Full structured summary
  try {
    const conversation = formatMessagesForSummary(messages);
    const prompt = buildFullSummaryPrompt(conversation, config);
    const result = await callback(prompt, maxOutputTokens);

    if (result != null && result.length > 0) {
      if (
        config?.summaryBudget != null &&
        config.summaryBudget > 0 &&
        config.tokenCounter != null
      ) {
        if (
          validateSummarySize(result, config.summaryBudget, config.tokenCounter)
        ) {
          return {
            summary: result,
            tier: 'full',
            messagesCompacted: messageCount,
          };
        }
      } else {
        return {
          summary: result,
          tier: 'full',
          messagesCompacted: messageCount,
        };
      }
    }
  } catch {
    // Fall through to Tier 2
  }

  // Tier 2: Simple summary
  try {
    const conversation = formatMessagesForSummary(messages);
    const prompt = buildSimpleSummaryPrompt(conversation);
    const result = await callback(prompt, SIMPLE_MAX_TOKENS);

    if (result != null && result.length > 0) {
      return {
        summary: result,
        tier: 'simple',
        messagesCompacted: messageCount,
      };
    }
  } catch {
    // Fall through to Tier 3
  }

  // Tier 3: Emergency (no LLM)
  const summary = createEmergencySummary(messages);
  return { summary, tier: 'emergency', messagesCompacted: messageCount };
}