import type {
  CodingRuntimeErrorCode,
  ConversationContentBlock,
  ConversationMessageNode,
  ConversationNode,
  ConversationSnapshot,
  ConversationToolNode,
  PublicUsage,
} from '../contracts';
import type {
  PiEventProjectorOptions,
  PiImageProjectionInput,
  PiProjectedAttachment,
} from './event-projector';
import { subagentDetailsOfResult } from '../subagent-protocol';

export interface PiSessionSnapshotInput {
  snapshot: ConversationSnapshot;
  workerGeneration: number;
  state: unknown;
  entries: unknown;
  stats?: unknown;
  projectImage?(image: PiImageProjectionInput): Promise<PiProjectedAttachment>;
}

interface SessionEntryRecord extends Record<string, unknown> {
  id: string;
  parentId: string | null;
  type: string;
}

export class PiSessionProjectionError extends Error {
  readonly code: CodingRuntimeErrorCode = 'CODING_SESSION_UNREADABLE';
  readonly recoverable = true;

  constructor() {
    super('Pi session data is unreadable; the source session was left unchanged.');
    this.name = 'PiSessionProjectionError';
  }
}

function unreadable(): never {
  throw new PiSessionProjectionError();
}

function asRecord(value: unknown): Record<string, unknown> | null {
  return value !== null && typeof value === 'object' && !Array.isArray(value)
    ? value as Record<string, unknown>
    : null;
}

function sessionEntry(value: unknown): SessionEntryRecord | null {
  const entry = asRecord(value);
  if (!entry
    || typeof entry.id !== 'string'
    || (entry.parentId !== null && typeof entry.parentId !== 'string')
    || typeof entry.type !== 'string') return null;
  return entry as SessionEntryRecord;
}

function usageOf(value: unknown): PublicUsage | undefined {
  const usage = asRecord(value);
  if (!usage
    || typeof usage.input !== 'number'
    || typeof usage.output !== 'number') return undefined;
  return {
    inputTokens: usage.input,
    outputTokens: usage.output,
    ...(typeof usage.cacheRead === 'number' ? { cacheReadTokens: usage.cacheRead } : {}),
    ...(typeof usage.cacheWrite === 'number' ? { cacheWriteTokens: usage.cacheWrite } : {}),
  };
}

function stopReasonOf(value: unknown): ConversationMessageNode['stopReason'] | undefined {
  if (value === 'toolUse') return 'tool-use';
  if (value === 'stop' || value === 'length' || value === 'error' || value === 'aborted') {
    return value;
  }
  return undefined;
}

function activePath(entriesValue: unknown, leafId: unknown): SessionEntryRecord[] {
  if (!Array.isArray(entriesValue)) unreadable();
  if (leafId !== null && typeof leafId !== 'string') unreadable();
  const index = new Map<string, SessionEntryRecord>();
  for (const value of entriesValue) {
    const entry = sessionEntry(value);
    if (!entry || index.has(entry.id)) unreadable();
    index.set(entry.id, entry);
  }
  if (leafId === null) return [];
  const path: SessionEntryRecord[] = [];
  const visited = new Set<string>();
  let current = index.get(leafId);
  if (!current) unreadable();
  while (current) {
    if (visited.has(current.id)) unreadable();
    visited.add(current.id);
    path.push(current);
    if (current.parentId === null) break;
    const parent = index.get(current.parentId);
    if (!parent) unreadable();
    current = parent;
  }
  return path.reverse();
}

function retainedTail(path: SessionEntryRecord[]): SessionEntryRecord[] {
  const compactionIndex = path.findLastIndex((entry) => entry.type === 'compaction');
  if (compactionIndex < 0) return path;
  const compaction = path[compactionIndex];
  const firstKeptEntryId = compaction.firstKeptEntryId;
  if (typeof firstKeptEntryId !== 'string') unreadable();
  const firstKeptIndex = path.findIndex(
    (entry, index) => index < compactionIndex && entry.id === firstKeptEntryId,
  );
  if (firstKeptIndex < 0) unreadable();
  return [
    compaction,
    ...path.slice(firstKeptIndex, compactionIndex),
    ...path.slice(compactionIndex + 1),
  ];
}

async function contentBlocks(
  messageId: string,
  content: unknown,
  conversationId: string,
  projectImage: PiEventProjectorOptions['projectImage'],
  idPart: 'content' | 'output',
): Promise<ConversationContentBlock[]> {
  const values = typeof content === 'string'
    ? [{ type: 'text', text: content }]
    : Array.isArray(content) ? content : [];
  const blocks: ConversationContentBlock[] = [];
  for (let contentIndex = 0; contentIndex < values.length; contentIndex += 1) {
    const block = asRecord(values[contentIndex]);
    if (block?.type === 'thinking' && typeof block.thinking === 'string') {
      blocks.push({
        kind: 'thinking',
        id: `${messageId}:${idPart}:${contentIndex}`,
        text: block.thinking,
        status: 'complete',
      });
    } else if (block?.type === 'text' && typeof block.text === 'string') {
      blocks.push({
        kind: 'text',
        id: `${messageId}:${idPart}:${contentIndex}`,
        text: block.text,
        status: 'complete',
      });
    } else if (block?.type === 'image'
      && typeof block.data === 'string'
      && typeof block.mimeType === 'string'
      && projectImage) {
      const attachment = await projectImage({
        conversationId,
        data: block.data,
        mime: block.mimeType,
        source: 'session',
      });
      blocks.push({
        kind: 'image',
        id: `${messageId}:${idPart}:${contentIndex}`,
        attachmentId: attachment.attachmentId,
        mime: attachment.mime,
      });
    }
  }
  return blocks;
}

function inputTextOf(value: unknown): string {
  if (typeof value === 'string') return value;
  try {
    return JSON.stringify(value ?? {}) ?? '{}';
  } catch {
    return '{}';
  }
}

async function projectEntries(
  path: SessionEntryRecord[],
  input: PiSessionSnapshotInput,
): Promise<ConversationNode[]> {
  const nodes: ConversationNode[] = [];
  const tools = new Map<string, ConversationToolNode>();
  const subagentIds = new Set<string>();
  for (const entry of path) {
    if (entry.type === 'compaction') {
      nodes.push({
        kind: 'compaction',
        id: `entry:${entry.id}`,
        runId: `session:${entry.id}`,
        source: 'automatic',
        status: 'complete',
        willRetry: false,
      });
      continue;
    }
    if (entry.type !== 'message') continue;
    const message = asRecord(entry.message);
    if (!message) unreadable();
    if (message.role === 'toolResult') {
      if (typeof message.toolCallId !== 'string') unreadable();
      const tool = tools.get(message.toolCallId);
      if (!tool) continue;
      tool.output = await contentBlocks(
        tool.id,
        message.content,
        input.snapshot.conversation.id,
        input.projectImage,
        'output',
      );
      tool.status = message.isError === true ? 'error' : 'complete';
      const details = subagentDetailsOfResult(message);
      if (details) {
        tool.details = details;
        const id = `subagent:${details.dispatchId}`;
        if (!subagentIds.has(id)) {
          subagentIds.add(id);
          nodes.push({
            kind: 'subagent',
            id,
            runId: input.snapshot.run.runId ?? `session:${entry.id}`,
            details,
          });
        }
      } else if (tool.toolName === 'subagent' && message.details !== undefined) {
        tool.output = [{
          kind: 'text',
          id: `${tool.id}:output:unavailable`,
          text: 'Subagent details are unavailable for this version.',
          status: 'complete',
        }];
      }
      continue;
    }
    if (message.role !== 'user' && message.role !== 'assistant') continue;
    const messageId = `entry:${entry.id}`;
    const usage = usageOf(message.usage);
    const stopReason = stopReasonOf(message.stopReason);
    const node: ConversationMessageNode = {
      kind: 'message',
      id: messageId,
      sourceEntryId: entry.id,
      role: message.role,
      status: message.stopReason === 'aborted'
        ? 'aborted'
        : message.stopReason === 'error' ? 'error' : 'complete',
      blocks: await contentBlocks(
        messageId,
        message.content,
        input.snapshot.conversation.id,
        input.projectImage,
        'content',
      ),
      ...(usage ? { usage } : {}),
      ...(stopReason ? { stopReason } : {}),
    };
    nodes.push(node);
    if (!Array.isArray(message.content)) continue;
    for (const value of message.content) {
      const block = asRecord(value);
      if (block?.type !== 'toolCall'
        || typeof block.id !== 'string'
        || typeof block.name !== 'string') continue;
      const tool: ConversationToolNode = {
        kind: 'tool',
        id: `tool:${block.id}`,
        toolCallId: block.id,
        toolName: block.name,
        title: block.name,
        inputText: inputTextOf(block.arguments),
        status: 'declared',
        output: [],
      };
      nodes.push(tool);
      tools.set(block.id, tool);
    }
  }
  return nodes;
}

function blockSignature(block: ConversationContentBlock): string {
  if (block.kind === 'image') return `image:${block.mime}`;
  return `${block.kind}:${block.text}`;
}

function messageSignature(node: ConversationMessageNode): string {
  return `${node.role}|${node.blocks.map(blockSignature).join('|')}`;
}

function rebaseBlocks(
  blocks: ConversationContentBlock[],
  nodeId: string,
  idPart: 'content' | 'output',
): ConversationContentBlock[] {
  return blocks.map((block, index) => ({ ...block, id: `${nodeId}:${idPart}:${index}` }));
}

function reconcileLiveIds(
  durableNodes: ConversationNode[],
  liveNodes: ConversationNode[],
): ConversationNode[] {
  const used = new Set<string>();
  return durableNodes.map((node) => {
    if (node.kind === 'message') {
      const live = liveNodes.find((candidate) => candidate.kind === 'message'
        && !used.has(candidate.id)
        && (candidate.sourceEntryId === node.sourceEntryId
          || messageSignature(candidate) === messageSignature(node)));
      if (!live || live.kind !== 'message') return node;
      used.add(live.id);
      return {
        ...node,
        id: live.id,
        blocks: rebaseBlocks(node.blocks, live.id, 'content'),
        ...(live.clientRequestId ? { clientRequestId: live.clientRequestId } : {}),
      };
    }
    if (node.kind === 'tool') {
      const live = liveNodes.find((candidate) => candidate.kind === 'tool'
        && !used.has(candidate.id)
        && candidate.toolCallId === node.toolCallId);
      if (!live || live.kind !== 'tool') return node;
      used.add(live.id);
      return {
        ...node,
        id: live.id,
        output: rebaseBlocks(node.output, live.id, 'output'),
      };
    }
    if (node.kind === 'compaction') {
      const live = liveNodes.findLast(
        (candidate) => candidate.kind === 'compaction' && !used.has(candidate.id),
      );
      if (!live || live.kind !== 'compaction') return node;
      used.add(live.id);
      return { ...node, id: live.id, runId: live.runId };
    }
    return node;
  });
}

function projectedContext(input: PiSessionSnapshotInput, state: Record<string, unknown>) {
  const stats = asRecord(input.stats);
  const contextUsage = asRecord(stats?.contextUsage);
  if (!contextUsage) {
    return {
      ...structuredClone(input.snapshot.context),
      compaction: state.isCompacting === true ? 'running' as const : 'idle' as const,
    };
  }
  const tokens = contextUsage.tokens;
  const recalculating = tokens === null;
  return {
    usedTokens: typeof tokens === 'number' ? tokens : 0,
    contextWindow: typeof contextUsage.contextWindow === 'number' ? contextUsage.contextWindow : 0,
    compaction: state.isCompacting === true ? 'running' as const : 'idle' as const,
    ...(recalculating ? { recalculating: true } : {}),
  };
}

export async function projectPiSessionSnapshot(
  input: PiSessionSnapshotInput,
): Promise<ConversationSnapshot> {
  if (!Number.isSafeInteger(input.workerGeneration) || input.workerGeneration <= 0) unreadable();
  const response = asRecord(input.entries);
  const state = asRecord(input.state);
  if (!response || !state) unreadable();
  const path = retainedTail(activePath(response.entries, response.leafId));
  const durableNodes = await projectEntries(path, input);
  const nodes = reconcileLiveIds(durableNodes, input.snapshot.nodes);
  return {
    ...structuredClone(input.snapshot),
    nodes,
    run: state.isStreaming === true
      ? { ...structuredClone(input.snapshot.run), status: 'running' }
      : { status: 'idle' },
    queue: { items: [] },
    context: projectedContext(input, state),
    pendingInteractions: [],
    worker: { status: 'ready', generation: input.workerGeneration },
    cursor: {
      workerGeneration: input.workerGeneration,
      seq: input.snapshot.cursor.workerGeneration === input.workerGeneration
        ? input.snapshot.cursor.seq
        : 0,
      ...(typeof response.leafId === 'string' ? { leafEntryId: response.leafId } : {}),
    },
  };
}