/**
 * @license
 * Copyright 2026 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import type { Content, Part } from '@google/genai';
import { randomUUID } from 'node:crypto';
import type {
  Episode,
  IrMetadata,
  SemanticPart,
  ToolExecution,
  AgentThought,
  AgentYield,
  UserPrompt,
} from './types.js';
import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';

// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references
const nodeIdentityMap = new WeakMap<object, string>();

export function getStableId(obj: object): string {
  let id = nodeIdentityMap.get(obj);
  if (!id) {
    id = randomUUID();
    nodeIdentityMap.set(obj, id);
  }
  return id;
}

function isRecord(v: unknown): v is Record<string, unknown> {
  return typeof v === 'object' && v !== null && !Array.isArray(v);
}

function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
  return (
    typeof ep.id === 'string' &&
    typeof ep.timestamp === 'number' &&
    Array.isArray(ep.concreteNodes) &&
    ep.concreteNodes.length > 0
  );
}

export function toIr(
  history: readonly Content[],
  tokenCalculator: ContextTokenCalculator,
): Episode[] {
  const episodes: Episode[] = [];
  let currentEpisode: Partial<Episode> | null = null;
  const pendingCallParts: Map<string, Part> = new Map();

  const createMetadata = (parts: Part[]): IrMetadata => {
    const tokens = tokenCalculator.estimateTokensForParts(parts, 0);
    return {
      originalTokens: tokens,
      currentTokens: tokens,
      transformations: [],
    };
  };

  const finalizeEpisode = () => {
    if (currentEpisode && isCompleteEpisode(currentEpisode)) {
      episodes.push(currentEpisode);
    }
    currentEpisode = null;
  };

  for (const msg of history) {
    if (!msg.parts) continue;

    if (msg.role === 'user') {
      const hasToolResponses = msg.parts.some((p) => !!p.functionResponse);
      const hasUserParts = msg.parts.some(
        (p) => !!p.text || !!p.inlineData || !!p.fileData,
      );

      if (hasToolResponses) {
        currentEpisode = parseToolResponses(
          msg,
          currentEpisode,
          pendingCallParts,
          tokenCalculator,
          createMetadata,
        );
      }

      if (hasUserParts) {
        finalizeEpisode();
        currentEpisode = parseUserParts(msg, createMetadata);
      }
    } else if (msg.role === 'model') {
      currentEpisode = parseModelParts(
        msg,
        currentEpisode,
        pendingCallParts,
        createMetadata,
      );
    }
  }

  if (currentEpisode) {
    finalizeYield(currentEpisode);
    finalizeEpisode();
  }

  return episodes;
}

function parseToolResponses(
  msg: Content,
  currentEpisode: Partial<Episode> | null,
  pendingCallParts: Map<string, Part>,
  tokenCalculator: ContextTokenCalculator,
  createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
  if (!currentEpisode) {
    currentEpisode = {
      id: getStableId(msg),
      timestamp: Date.now(),
      concreteNodes: [],
    };
  }

  for (const part of msg.parts!) {
    if (part.functionResponse) {
      const callId = part.functionResponse.id || '';
      const matchingCall = pendingCallParts.get(callId);

      const intentTokens = matchingCall
        ? tokenCalculator.estimateTokensForParts([matchingCall])
        : 0;
      const obsTokens = tokenCalculator.estimateTokensForParts([part]);

      const step: ToolExecution = {
        id: getStableId(part),
        type: 'TOOL_EXECUTION',
        toolName: part.functionResponse.name || 'unknown',
        intent: isRecord(matchingCall?.functionCall?.args)
          ? matchingCall.functionCall.args
          : {},
        observation: isRecord(part.functionResponse.response)
          ? part.functionResponse.response
          : {},
        tokens: {
          intent: intentTokens,
          observation: obsTokens,
        },
        metadata: {
          originalTokens: intentTokens + obsTokens,
          currentTokens: intentTokens + obsTokens,
          transformations: [],
        },
      };
      currentEpisode.concreteNodes = [
        ...(currentEpisode.concreteNodes || []),
        step,
      ];
      if (callId) pendingCallParts.delete(callId);
    }
  }
  return currentEpisode;
}

function parseUserParts(
  msg: Content,
  createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
  const semanticParts: SemanticPart[] = [];
  for (const p of msg.parts!) {
    if (p.text !== undefined)
      semanticParts.push({ type: 'text', text: p.text });
    else if (p.inlineData)
      semanticParts.push({
        type: 'inline_data',
        mimeType: p.inlineData.mimeType || '',
        data: p.inlineData.data || '',
      });
    else if (p.fileData)
      semanticParts.push({
        type: 'file_data',
        mimeType: p.fileData.mimeType || '',
        fileUri: p.fileData.fileUri || '',
      });
    else if (!p.functionResponse)
      semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns
  }

  const trigger: UserPrompt = {
    id: getStableId(msg.parts![0] || msg),
    type: 'USER_PROMPT',
    semanticParts,
    metadata: createMetadata(msg.parts!.filter((p) => !p.functionResponse)),
  };

  return {
    id: getStableId(msg),
    timestamp: Date.now(),
    concreteNodes: [trigger],
  };
}

function parseModelParts(
  msg: Content,
  currentEpisode: Partial<Episode> | null,
  pendingCallParts: Map<string, Part>,
  createMetadata: (parts: Part[]) => IrMetadata,
): Partial<Episode> {
  if (!currentEpisode) {
    currentEpisode = {
      id: getStableId(msg),
      timestamp: Date.now(),
      concreteNodes: [],
    };
  }

  for (const part of msg.parts!) {
    if (part.functionCall) {
      const callId = part.functionCall.id || '';
      if (callId) pendingCallParts.set(callId, part);
    } else if (part.text) {
      const thought: AgentThought = {
        id: getStableId(part),
        type: 'AGENT_THOUGHT',
        text: part.text,
        metadata: createMetadata([part]),
      };
      currentEpisode.concreteNodes = [
        ...(currentEpisode.concreteNodes || []),
        thought,
      ];
    }
  }
  return currentEpisode;
}

function finalizeYield(currentEpisode: Partial<Episode>) {
  if (currentEpisode.concreteNodes && currentEpisode.concreteNodes.length > 0) {
    const yieldNode: AgentYield = {
      id: randomUUID(),
      type: 'AGENT_YIELD',
      text: 'Yield', // Synthesized yield since we don't have the original concrete node
      metadata: {
        originalTokens: 1,
        currentTokens: 1,
        transformations: [],
      },
    };
    const existingNodes =
      currentEpisode.concreteNodes as Array<import('./types.js').ConcreteNode>;
    currentEpisode.concreteNodes = [...existingNodes, yieldNode];
  }
}