import { PersonaMessage } from './types';
import { FileMessagePart, ImageMessagePart, ThreadMessageLike, ThreadUserMessagePart, MessageStatus } from '@assistant-ui/react';

function sanitizeToolResult(result: unknown): unknown {
  if (result === null || result === undefined) return result;
  if (typeof result === 'string') return result;
  if (Array.isArray(result)) return result.map(sanitizeToolResult);
  if (typeof result === 'object') {
    const r = result as Record<string, unknown>;
    // Image-shaped: { content, contentType }. Keep the pixels: the UI renders them
    // (see ToolFallback/MediaResult). Only live messages ever get here with `content`
    // set — core marks Image.content transient, so messages reloaded from history
    // arrive without it and fall through to the generic branch below.
    if (typeof r['contentType'] === 'string' && typeof r['content'] === 'string') {
      return r;
    }
    // File-shaped: has name + contentType (strip base64, keep metadata)
    if (typeof r['name'] === 'string' && typeof r['contentType'] === 'string') {
      return { type: 'file', name: r['name'], url: r['url'] ?? null, contentType: r['contentType'] };
    }
    return Object.fromEntries(Object.entries(r).map(([k, v]) => [k, sanitizeToolResult(v)]));
  }
  return result;
}

/** Messaggio marcato dal runtime come avviso di servizio, non risposta dell'agente. */
function isNotice(message: PersonaMessage): boolean {
  return message.metadata?.custom?.notice !== undefined;
}

let messageIdCounter = 0;

export function generateMessageId(): string {
  return `msg_${Date.now()}_${messageIdCounter++}`;
}

export function ensureMessageId(message: PersonaMessage): PersonaMessage {
  // Convert thought to reasoning type
  if (message.thought) {
    message = { ...message, type: 'reasoning' };
  }

  if (!message.id) {
    return {
      ...message,
      id: generateMessageId(),
      createdAt: message.createdAt || new Date(),
    };
  }
  if (!message.createdAt) {
    return { ...message, createdAt: new Date() };
  }
  return message;
}

function removeEmptyMessages(messages: PersonaMessage[]): PersonaMessage[] {
  return messages.filter((message) => {
    if (message.role === 'assistant' && !message.functionCalls && !message.file && !message.image) {
      if (!message.text || message.text.trim() === '') {
        return false;
      }
    }
    return true;
  });
}
function parseMessages(messages: PersonaMessage[]): PersonaMessage[] {
  const outputMessages: PersonaMessage[] = [];
  let currentMessage: PersonaMessage | null = null;

  for (let message of messages) {
    // Ensure every message has an ID
    message = ensureMessageId(message);

    if (message.type === 'transaction') {
      continue;
    }

    if (message.file) {
      if (currentMessage) {
        outputMessages.push(currentMessage);
        currentMessage = null;
      }
      outputMessages.push(message);
      continue;
    } else if (message.functionCalls) {
      if (currentMessage) {
        // Drop empty assistant messages that precede function calls (streaming artifacts)
        const isEmptyAssistant =
          currentMessage.role === 'assistant' &&
          (!currentMessage.text || currentMessage.text.trim() === '') &&
          !currentMessage.file &&
          !currentMessage.image;
        if (!isEmptyAssistant) {
          outputMessages.push(currentMessage);
        }
      }
      outputMessages.push(message);
      currentMessage = null;
    } else if (message.functionResponse) {
      outputMessages[outputMessages.length - 1] = {
        ...outputMessages[outputMessages.length - 1],
        functionResponse: message.functionResponse,
      };
    } else if (
      currentMessage &&
      message.type === currentMessage.type &&
      message.protocol === currentMessage.protocol &&
      currentMessage.role === message.role &&
      // Un avviso di sistema (turno interrotto, errore) non è uno chunk di
      // streaming: se lo si fonde finisce appiccicato in fondo alla risposta
      // troncata dell'agente.
      !isNotice(message) &&
      !isNotice(currentMessage) &&
      // CRITICAL: Only merge if the new message has the same ID (streaming chunks)
      // OR if the current message doesn't have a finishReason yet (still streaming)
      // This preserves branches while allowing streaming to work
      (message.id === currentMessage.id || !currentMessage.finishReason)
    ) {
      // Merge streaming chunks
      currentMessage.text += message.text;
      currentMessage.file = currentMessage.file || message.file;
      currentMessage.attachments = currentMessage.attachments?.length
        ? currentMessage.attachments
        : message.attachments;
      currentMessage.finishReason = message.finishReason || currentMessage.finishReason;
      // Keep the ID of the first chunk
      currentMessage.id = currentMessage.id || message.id;
      currentMessage.sources = currentMessage.sources ? [...(currentMessage.sources || []), ...(message.sources || [])] : message.sources;
    } else {
      // Different message or current is complete - don't merge
      if (currentMessage) {
        outputMessages.push(currentMessage);
      }
      currentMessage = {
        ...message,
        text: message.file ? '' : message.text,
        finishReason: message.file ? undefined : message.finishReason,
      };
    }
  }

  if (currentMessage) {
    outputMessages.push(currentMessage);
  }
  const cleanMessages = removeEmptyMessages(outputMessages);
  return cleanMessages;
}

function convertMessage(message: PersonaMessage): ThreadMessageLike {
  const files = message.file
    ? [
        {
          type: 'file',
          data: message.file.url,
          mimeType: message.file.contentType,
        } as FileMessagePart,
      ]
    : [];

  // Allegati del messaggio utente: l'eco locale dell'invio e i messaggi
  // ricaricati dallo storico diventano image/file part così restano visibili
  // nella bolla del messaggio (le immagini come data URL).
  const attachmentParts: ThreadUserMessagePart[] = [];
  for (const att of message.attachments ?? []) {
    if (att.image) {
      attachmentParts.push({
        type: 'image',
        image: `data:${att.image.contentType};base64,${att.image.content}`,
      } as ImageMessagePart);
    }
    if (att.file) {
      const data = att.file.url ?? (att.file.base64 ? `data:${att.file.contentType};base64,${att.file.base64}` : undefined);
      if (data) {
        attachmentParts.push({
          type: 'file',
          data,
          mimeType: att.file.contentType,
          filename: att.file.name,
        } as FileMessagePart);
      }
    }
  }
  // Retrocompatibilità: messaggi utente con `image` (p.e. dallo storico del
  // server) — stesso trattamento dell'eco.
  if (message.role === 'user' && message.image) {
    attachmentParts.push({
      type: 'image',
      image: `data:${message.image.contentType};base64,${message.image.content}`,
    } as ImageMessagePart);
  }

  // Determine message status for assistant messages
  let messageStatus: MessageStatus | undefined;

  if (message.role === 'assistant' || message.role === 'function') {
    if (message.status) {
      // Use existing status if provided
      messageStatus = message.status as MessageStatus;
    } else if (message.finishReason === 'stop') {
      messageStatus = { type: 'complete', reason: 'stop' } as const;
    } else if (message.finishReason === 'function_call') {
      messageStatus = { type: 'complete', reason: 'stop' } as const;
    } else if (!message.finishReason) {
      // Message is still being generated - this will trigger optimistic updates
      messageStatus = undefined; // Let the runtime handle in-progress state
    }
  }

  // Handle function/tool calls
  if (message.role === 'function') {
    const result: ThreadMessageLike = {
      id: message.id!,
      role: 'assistant',
      createdAt: message.createdAt,
      content:
        message.functionCalls?.map((call) => ({
          type: 'tool-call',
          toolName: call.name,
          toolCallId: call.id || `tool_${Date.now()}`,
          args: call.args,
          result: sanitizeToolResult(message.functionResponse?.result),
        })) ?? [],
      metadata: {
        ...message.metadata,
        custom: {
          protocol: message.protocol,
          sessionId: message.sessionId,
          sources: message.sources,
          sender: message.sender ?? undefined,
          ...message.metadata?.custom,
        },
      },
    };

    if (messageStatus) {
      return { ...result, status: messageStatus };
    }
    return result;
  }

  // Handle regular messages
  const baseMessage: ThreadMessageLike = {
    id: message.id!,
    role: message.role,
    createdAt: message.createdAt,
    content:
      message.type === 'reasoning'
        ? [{ type: 'reasoning' as 'text', text: message.text }, ...files]
        : [{ type: 'text', text: message.text }, ...attachmentParts, ...files],
    metadata: {
      ...message.metadata,
      custom: {
        protocol: message.protocol,
        sessionId: message.sessionId,
        sources: message.sources,
        sender: message.sender ?? undefined,
        type: message.type,
        ...message.metadata?.custom,
      },
    },
  };

  // Add status only for assistant messages
  if (message.role === 'assistant' && messageStatus) {
    return { ...baseMessage, status: messageStatus };
  }

  return baseMessage;
}

function isNoticeLike(message: ThreadMessageLike): boolean {
  const custom = message.metadata?.custom as Record<string, unknown> | undefined;
  return custom?.notice !== undefined;
}

function isTurnPart(message: ThreadMessageLike): boolean {
  return message.role === 'assistant' && !isNoticeLike(message);
}

/**
 * Builds a converter that folds every run of consecutive agent messages
 * (reasoning, tool calls, text) into a single assistant message per turn, the
 * shape assistant-ui expects: grouping, action bars and "N tool calls" all work
 * per message. User messages and runtime notices stay on their own.
 *
 * The converter is stateful on purpose: conversions are cached per source
 * message and per turn, so while streaming only the last turn is rebuilt and
 * finished turns keep their identity (no re-render, no tool UI remount).
 */
function createTurnGrouper() {
  const converted = new WeakMap<PersonaMessage, ThreadMessageLike>();
  let turnCache = new Map<string, { members: ThreadMessageLike[]; turn: ThreadMessageLike }>();

  const convert = (message: PersonaMessage) => {
    let result = converted.get(message);
    if (!result) {
      result = convertMessage(message);
      converted.set(message, result);
    }
    return result;
  };

  const buildTurn = (members: ThreadMessageLike[]): ThreadMessageLike => {
    const [first] = members;
    const last = members[members.length - 1];
    const partsOf = (m: ThreadMessageLike) =>
      typeof m.content === 'string' ? [{ type: 'text' as const, text: m.content }] : m.content;
    const content = members.flatMap(partsOf);
    // When each part arrived (its source message's createdAt, epoch ms): once
    // merged, the turn has a single createdAt, and the UI needs these to tell
    // how long the agent worked ("Thought for 12s").
    const partTimestamps = members.flatMap((m) => {
      const time = m.createdAt ? new Date(m.createdAt).getTime() : NaN;
      return partsOf(m).map(() => (Number.isFinite(time) ? time : null));
    });
    const sources = members.flatMap((m) => ((m.metadata?.custom as { sources?: unknown[] } | undefined)?.sources ?? []));
    // Who answered: the agent behind the last message that says so (the
    // reasoning packets carry no sender, and after a handoff the turn ends
    // with the agent that took over).
    const sender = members
      .map((m) => (m.metadata?.custom as { sender?: string } | undefined)?.sender)
      .filter(Boolean)
      .pop();
    // The turn is as far along as its last message: a still-streaming tail has
    // no status, and must not inherit the "complete" of the earlier messages.
    const { status: _firstStatus, ...rest } = first;
    return {
      ...rest,
      content,
      ...(last.status ? { status: last.status } : {}),
      metadata: {
        ...first.metadata,
        custom: { ...first.metadata?.custom, sources: sources.length ? sources : undefined, sender, partTimestamps },
      },
    } as ThreadMessageLike;
  };

  return (messages: readonly PersonaMessage[]): ThreadMessageLike[] => {
    const output: ThreadMessageLike[] = [];
    const nextCache: typeof turnCache = new Map();
    let members: ThreadMessageLike[] = [];

    const flush = () => {
      if (members.length === 0) return;
      const key = String(members[0].id);
      const cached = turnCache.get(key);
      const turn =
        cached && cached.members.length === members.length && cached.members.every((m, i) => m === members[i])
          ? cached.turn
          : buildTurn(members);
      nextCache.set(key, { members, turn });
      output.push(turn);
      members = [];
    };

    for (const message of messages) {
      const item = convert(message);
      if (isTurnPart(item)) {
        members.push(item);
      } else {
        flush();
        output.push(item);
      }
    }
    flush();
    turnCache = nextCache;
    return output;
  };
}

export { parseMessages, convertMessage, removeEmptyMessages, createTurnGrouper };