import type {
  SessionTreeCustomRuntimeMessage,
  SessionTreeEntry,
  SessionTreeEntryIndex,
  SessionTreeRuntimeContext,
  SessionTreeRuntimeMessage,
} from './types.js';

export function buildSessionContext(
  entries: SessionTreeEntry[],
  leafId?: string | null,
  byId?: SessionTreeEntryIndex,
): SessionTreeRuntimeContext {
  if (leafId === null) {
    return {
      messages: [],
      thinkingLevel: 'off',
      model: null,
    };
  }

  const index = createEntryLookup(entries, byId);
  const leaf = resolveLeaf(entries, leafId, index);

  if (!leaf) {
    return {
      messages: [],
      thinkingLevel: 'off',
      model: null,
    };
  }

  const path = getPathToLeafFromIndex(leaf.id, index);

  return buildSessionContextFromPath(path);
}

export function buildSessionContextFromPath(path: SessionTreeEntry[]): SessionTreeRuntimeContext {
  if (path.length === 0) {
    return {
      messages: [],
      thinkingLevel: 'off',
      model: null,
    };
  }

  let thinkingLevel = 'off';
  let model: SessionTreeRuntimeContext['model'] = null;
  let latestCompactionIndex = -1;

  for (let indexOnPath = 0; indexOnPath < path.length; indexOnPath += 1) {
    const entry = path[indexOnPath];

    if (entry.type === 'thinking_level_change') {
      thinkingLevel = entry.thinkingLevel;
      continue;
    }

    if (entry.type === 'model_change') {
      model = {
        provider: entry.provider,
        modelId: entry.modelId,
      };
      continue;
    }

    if (entry.type === 'message' && entry.message.role === 'assistant') {
      const provider = extractAssistantProvider(entry.message);
      const modelId = extractAssistantModelId(entry.message);

      if (provider && modelId) {
        model = {
          provider,
          modelId,
        };
      }
      continue;
    }

    if (entry.type === 'compaction') {
      latestCompactionIndex = indexOnPath;
    }
  }

  if (latestCompactionIndex >= 0) {
    const compaction = path[latestCompactionIndex];
    if (compaction?.type !== 'compaction') {
      throw new Error('Latest compaction index did not resolve to a compaction entry');
    }

    const messages: SessionTreeRuntimeMessage[] = [toCompactionSummaryMessage(compaction)];
    const firstKeptIndex = path.findIndex(
      (entry, indexOnPath) =>
        indexOnPath < latestCompactionIndex && entry.id === compaction.firstKeptEntryId,
    );

    if (firstKeptIndex < 0) {
      throw new Error(
        `Compaction ${compaction.id} firstKeptEntryId ${compaction.firstKeptEntryId} not found on active path`,
      );
    }

    for (let indexOnPath = firstKeptIndex; indexOnPath < latestCompactionIndex; indexOnPath += 1) {
      appendRuntimeMessage(messages, path[indexOnPath]);
    }

    for (let indexOnPath = latestCompactionIndex + 1; indexOnPath < path.length; indexOnPath += 1) {
      appendRuntimeMessage(messages, path[indexOnPath]);
    }

    return {
      messages,
      thinkingLevel,
      model,
    };
  }

  const messages: SessionTreeRuntimeMessage[] = [];
  for (const entry of path) {
    appendRuntimeMessage(messages, entry);
  }

  return {
    messages,
    thinkingLevel,
    model,
  };
}

interface EntryLookup {
  get(id: string): SessionTreeEntry | undefined;
}

function createEntryLookup(entries: SessionTreeEntry[], byId: SessionTreeEntryIndex | undefined): EntryLookup {
  const index = buildEntryIndex(entries);

  if (byId instanceof Map) {
    for (const [id, entry] of byId) {
      index[id] = entry;
    }

    return {
      get(id: string): SessionTreeEntry | undefined {
        return index[id];
      },
    };
  }

  if (byId) {
    for (const id of Object.keys(byId)) {
      index[id] = byId[id];
    }
  }

  return {
    get(id: string): SessionTreeEntry | undefined {
      if (!Object.prototype.hasOwnProperty.call(index, id)) {
        return undefined;
      }

      return index[id];
    },
  };
}

function buildEntryIndex(entries: SessionTreeEntry[]): Record<string, SessionTreeEntry> {
  return entries.reduce<Record<string, SessionTreeEntry>>((index, entry) => {
    index[entry.id] = entry;
    return index;
  }, Object.create(null) as Record<string, SessionTreeEntry>);
}

function getPathToLeafFromIndex(leafEntryId: string, index: EntryLookup): SessionTreeEntry[] {
  const leaf = index.get(leafEntryId);

  if (!leaf) {
    throw new Error(`Entry ${leafEntryId} not found`);
  }

  const path: SessionTreeEntry[] = [];
  const visited = new Set<string>();
  let current: SessionTreeEntry | undefined = leaf;

  while (current) {
    if (visited.has(current.id)) {
      throw new Error('Cycle detected');
    }

    visited.add(current.id);
    path.push(current);

    if (current.parentId === null) {
      break;
    }

    const parent = index.get(current.parentId);
    if (!parent) {
      throw new Error(`Parent entry ${current.parentId} not found`);
    }

    current = parent;
  }

  return path.reverse();
}

function resolveLeaf(
  entries: SessionTreeEntry[],
  leafId: string | null | undefined,
  byId: EntryLookup,
): SessionTreeEntry | undefined {
  if (leafId !== undefined && leafId !== null) {
    return byId.get(leafId) ?? entries[entries.length - 1];
  }

  return entries[entries.length - 1];
}

function appendRuntimeMessage(
  messages: SessionTreeRuntimeMessage[],
  entry: SessionTreeEntry | undefined,
): void {
  if (!entry) {
    return;
  }

  if (entry.type === 'message') {
    messages.push(entry.message);
    return;
  }

  if (entry.type === 'custom_message') {
    messages.push(toCustomRuntimeMessage(entry));
    return;
  }

  if (entry.type === 'branch_summary') {
    messages.push({
      role: 'branchSummary',
      summary: entry.summary,
      fromId: entry.fromId,
      timestamp: new Date(entry.timestamp).getTime(),
    });
  }
}

function toCustomRuntimeMessage(entry: Extract<SessionTreeEntry, { type: 'custom_message' }>): SessionTreeCustomRuntimeMessage {
  return {
    role: 'custom',
    customType: entry.customType,
    content: entry.content,
    display: entry.display,
    details: entry.details,
    timestamp: new Date(entry.timestamp).getTime(),
  };
}

function toCompactionSummaryMessage(entry: Extract<SessionTreeEntry, { type: 'compaction' }>): SessionTreeRuntimeMessage {
  return {
    role: 'compactionSummary',
    summary: entry.summary,
    tokensBefore: entry.tokensBefore,
    timestamp: new Date(entry.timestamp).getTime(),
  };
}

function extractAssistantProvider(message: Record<string, unknown>): string | null {
  const directProvider = message.provider;
  if (typeof directProvider === 'string' && directProvider.length > 0) {
    return directProvider;
  }

  const nestedProvider = message.metadata;
  if (
    nestedProvider &&
    typeof nestedProvider === 'object' &&
    typeof (nestedProvider as Record<string, unknown>).provider === 'string'
  ) {
    return (nestedProvider as Record<string, string>).provider;
  }

  return null;
}

function extractAssistantModelId(message: Record<string, unknown>): string | null {
  const directModel = message.model;
  if (typeof directModel === 'string' && directModel.length > 0) {
    return directModel;
  }

  return null;
}
