import * as fs from 'fs';
import * as readline from 'readline';

export interface ParsedMessage {
  id: string;
  parentId: string | null;
  timestamp: string;
  role: 'user' | 'assistant';
  content: string | null;
  contentType: 'text' | 'tool_call' | 'thinking' | 'tool_result';
  toolName: string | null;
  tokens: number | null;
  inputTokens: number | null;
  outputTokens: number | null;
  totalTokens: number | null;
}

export interface SessionHeader {
  id: string;
  timestamp: string;
  version: number;
}

function extractTextContent(contentArr: any[]): string | null {
  if (!Array.isArray(contentArr)) return typeof contentArr === 'string' ? contentArr : null;

  const parts: string[] = [];
  for (const block of contentArr) {
    if (block.type === 'text' && block.text) {
      parts.push(block.text);
    }
  }
  return parts.length > 0 ? parts.join('\n') : null;
}

function extractToolCalls(contentArr: any[]): { name: string; content: string }[] {
  if (!Array.isArray(contentArr)) return [];
  return contentArr
    .filter((b: any) => b.type === 'toolCall')
    .map((b: any) => ({
      name: b.name || 'unknown',
      content: JSON.stringify(b.arguments || {}).slice(0, 2000),
    }));
}

function extractThinking(contentArr: any[]): string | null {
  if (!Array.isArray(contentArr)) return null;
  const thinking = contentArr.find((b: any) => b.type === 'thinking');
  return thinking?.thinking || null;
}

export async function parseJSONLFile(
  filePath: string,
  startByte: number = 0
): Promise<{ messages: ParsedMessage[]; header: SessionHeader | null; bytesRead: number; linesRead: number }> {
  const messages: ParsedMessage[] = [];
  let header: SessionHeader | null = null;
  let linesRead = 0;

  const stat = fs.statSync(filePath);
  if (stat.size <= startByte) {
    return { messages, header, bytesRead: startByte, linesRead: 0 };
  }

  const stream = fs.createReadStream(filePath, {
    start: startByte,
    encoding: 'utf-8',
  });

  const rl = readline.createInterface({ input: stream, crlfDelay: Infinity });

  for await (const line of rl) {
    linesRead++;
    if (!line.trim()) continue;

    let entry: any;
    try {
      entry = JSON.parse(line);
    } catch {
      continue;
    }

    // Session header
    if (entry.type === 'session') {
      header = {
        id: entry.id,
        timestamp: entry.timestamp,
        version: entry.version || 1,
      };
      continue;
    }

    // Message entries
    if (entry.type === 'message' && entry.message) {
      const msg = entry.message;
      const role = msg.role;
      if (role !== 'user' && role !== 'assistant') continue;

      const content = extractTextContent(msg.content);
      const toolCalls = extractToolCalls(msg.content);
      const thinking = extractThinking(msg.content);

      // Main text message
      if (content) {
        messages.push({
          id: entry.id,
          parentId: entry.parentId || null,
          timestamp: entry.timestamp,
          role,
          content,
          contentType: 'text',
          toolName: null,
          tokens: null,
          inputTokens: msg.usage?.input || null,
          outputTokens: msg.usage?.output || null,
          totalTokens: msg.usage?.totalTokens || null,
        });
      }

      // Tool calls as separate entries
      for (let i = 0; i < toolCalls.length; i++) {
        messages.push({
          id: `${entry.id}_tool_${i}`,
          parentId: entry.id,
          timestamp: entry.timestamp,
          role: 'assistant',
          content: toolCalls[i].content,
          contentType: 'tool_call',
          toolName: toolCalls[i].name,
          tokens: null,
          inputTokens: null,
          outputTokens: null,
          totalTokens: null,
        });
      }

      // Thinking block
      if (thinking && !content) {
        messages.push({
          id: `${entry.id}_think`,
          parentId: entry.id,
          timestamp: entry.timestamp,
          role: 'assistant',
          content: thinking.slice(0, 5000),
          contentType: 'thinking',
          toolName: null,
          tokens: null,
          inputTokens: msg.usage?.input || null,
          outputTokens: msg.usage?.output || null,
          totalTokens: msg.usage?.totalTokens || null,
        });
      }

      // If assistant message has no text but has usage, still record it as a token counter
      if (role === 'assistant' && !content && toolCalls.length === 0 && !thinking && msg.usage) {
        // Skip — no meaningful content to record
      }
    }

    // Tool result entries
    if (entry.type === 'toolResult') {
      const resultContent = typeof entry.content === 'string'
        ? entry.content.slice(0, 3000)
        : JSON.stringify(entry.content || '').slice(0, 3000);

      messages.push({
        id: entry.id,
        parentId: entry.parentId || null,
        timestamp: entry.timestamp,
        role: 'assistant',
        content: resultContent,
        contentType: 'tool_result',
        toolName: entry.toolName || null,
        tokens: null,
        inputTokens: null,
        outputTokens: null,
        totalTokens: null,
      });
    }
  }

  const bytesRead = stat.size;
  return { messages, header, bytesRead, linesRead };
}

export interface SessionIndex {
  sessionId: string;
  chatType: string;
  channel: string;
  origin: {
    provider?: string;
    surface?: string;
    chatType?: string;
    label?: string;
    from?: string;
    to?: string;
    accountId?: string;
  };
  sessionFile: string;
  inputTokens: number;
  outputTokens: number;
  totalTokens: number;
  model: string;
  updatedAt: number;
}

export function parseSessionsIndex(filePath: string): Map<string, SessionIndex> {
  const map = new Map<string, SessionIndex>();
  try {
    const raw = JSON.parse(fs.readFileSync(filePath, 'utf-8'));
    for (const [key, val] of Object.entries(raw)) {
      const v = val as any;
      map.set(key, {
        sessionId: v.sessionId,
        chatType: v.chatType || 'direct',
        channel: v.lastChannel || v.deliveryContext?.channel || 'telegram',
        origin: v.origin || {},
        sessionFile: v.sessionFile,
        inputTokens: v.inputTokens || 0,
        outputTokens: v.outputTokens || 0,
        totalTokens: v.totalTokens || 0,
        model: v.model || 'unknown',
        updatedAt: v.updatedAt || Date.now(),
      });
    }
  } catch (err) {
    console.error(`Failed to parse sessions index ${filePath}:`, err);
  }
  return map;
}