import { Injectable, Logger } from '@nestjs/common';
import { MessageRepository } from '../../infrastructure/repositories/message.repository';
import { ConversationMessage } from '../entities/conversation-message.entity';
import { TenantContextService } from '@it0/common';
import * as crypto from 'crypto';

export interface AnthropicHistoryMessage {
  role: 'user' | 'assistant';
  content: string | any[];
}

@Injectable()
export class ConversationContextService {
  private readonly logger = new Logger(ConversationContextService.name);

  constructor(private readonly messageRepository: MessageRepository) {}

  /**
   * Save a user message to the conversation history.
   */
  async saveUserMessage(
    sessionId: string,
    content: string,
    attachments?: Array<{ base64Data: string; mediaType: string; fileName?: string }>,
  ): Promise<ConversationMessage> {
    const tenantId = TenantContextService.getTenantId();
    const sequenceNumber = await this.messageRepository.getNextSequenceNumber(sessionId);

    const message = new ConversationMessage();
    message.id = crypto.randomUUID();
    message.tenantId = tenantId;
    message.sessionId = sessionId;
    message.role = 'user';
    message.content = content;
    message.attachments = attachments;
    message.tokenCount = this.estimateTokens(content);
    message.sequenceNumber = sequenceNumber;
    message.createdAt = new Date();

    return this.messageRepository.save(message);
  }

  /**
   * Save an assistant response to the conversation history.
   */
  async saveAssistantMessage(
    sessionId: string,
    content: string,
    toolCalls?: any[],
    toolResults?: any[],
  ): Promise<ConversationMessage> {
    const tenantId = TenantContextService.getTenantId();
    const sequenceNumber = await this.messageRepository.getNextSequenceNumber(sessionId);

    const message = new ConversationMessage();
    message.id = crypto.randomUUID();
    message.tenantId = tenantId;
    message.sessionId = sessionId;
    message.role = 'assistant';
    message.content = content;
    message.toolCalls = toolCalls;
    message.toolResults = toolResults;
    message.tokenCount = this.estimateTokens(content);
    message.sequenceNumber = sequenceNumber;
    message.createdAt = new Date();

    return this.messageRepository.save(message);
  }

  /**
   * Load conversation history for a session, formatted as Anthropic messages.
   * Returns the most recent `maxMessages` messages in chronological order.
   */
  async loadContext(
    sessionId: string,
    maxMessages = 20,
  ): Promise<AnthropicHistoryMessage[]> {
    const messages = await this.messageRepository.findRecentBySessionId(
      sessionId,
      maxMessages,
    );

    if (messages.length === 0) return [];

    const history: AnthropicHistoryMessage[] = [];

    for (const msg of messages) {
      if (msg.role === 'user') {
        if (msg.attachments && msg.attachments.length > 0) {
          // Build multimodal content blocks for messages with attachments
          const contentBlocks: any[] = [];
          for (const att of msg.attachments) {
            if (att.mediaType === 'application/pdf') {
              // PDF → Anthropic document content block
              contentBlocks.push({
                type: 'document',
                source: {
                  type: 'base64',
                  media_type: 'application/pdf',
                  data: att.base64Data,
                },
              });
            } else {
              // Image → Anthropic image content block
              contentBlocks.push({
                type: 'image',
                source: {
                  type: 'base64',
                  media_type: att.mediaType,
                  data: att.base64Data,
                },
              });
            }
          }
          if (msg.content && msg.content !== '[图片]') {
            contentBlocks.push({ type: 'text', text: msg.content });
          }
          history.push({ role: 'user', content: contentBlocks });
        } else {
          history.push({ role: 'user', content: msg.content });
        }
      } else if (msg.role === 'assistant') {
        // If the assistant message has tool calls, build content blocks
        if (msg.toolCalls && msg.toolCalls.length > 0) {
          const contentBlocks: any[] = [];
          if (msg.content) {
            contentBlocks.push({ type: 'text', text: msg.content });
          }
          for (const tc of msg.toolCalls) {
            contentBlocks.push({
              type: 'tool_use',
              id: tc.id,
              name: tc.name,
              input: tc.input,
            });
          }
          history.push({ role: 'assistant', content: contentBlocks });

          // Add tool results as the next user message
          if (msg.toolResults && msg.toolResults.length > 0) {
            history.push({
              role: 'user',
              content: msg.toolResults.map((tr: any) => ({
                type: 'tool_result',
                tool_use_id: tr.tool_use_id,
                content: tr.content,
              })),
            });
          }
        } else {
          history.push({ role: 'assistant', content: msg.content });
        }
      }
    }

    // Ensure messages alternate correctly — first message should be 'user'
    if (history.length > 0 && history[0].role !== 'user') {
      history.shift();
    }

    this.logger.log(
      `Loaded ${history.length} context messages for session ${sessionId}`,
    );

    return history;
  }

  /**
   * Build a text summary of conversation history for SDK engine prompt prefix.
   */
  async buildPromptPrefix(
    sessionId: string,
    maxMessages = 20,
  ): Promise<string> {
    const messages = await this.messageRepository.findRecentBySessionId(
      sessionId,
      maxMessages,
    );

    if (messages.length === 0) return '';

    const lines: string[] = ['[Previous conversation]'];
    for (const msg of messages) {
      const role = msg.role === 'user' ? 'User' : 'Assistant';
      // Truncate long messages in the prefix
      const content = msg.content.length > 500
        ? msg.content.slice(0, 500) + '...'
        : msg.content;
      lines.push(`${role}: ${content}`);
    }
    lines.push('', '[Current request]');

    return lines.join('\n');
  }

  /**
   * Simple token estimation: ~4 characters per token for English,
   * ~2 characters per token for CJK.
   */
  private estimateTokens(content: string): number {
    if (!content) return 0;
    // Count CJK characters
    const cjkCount = (content.match(/[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]/g) || []).length;
    const nonCjkCount = content.length - cjkCount;
    return Math.ceil(cjkCount / 2 + nonCjkCount / 4);
  }
}
