import ComposableArchitecture
import Foundation

private let JSON_FIELD_MESSAGE = "message"

/// Handles tool call events during message generation.
///
/// This handler processes tool call arguments as they stream in, extracting partial
/// message content from JSON arguments and updating messages accordingly.
public struct MessageToolCallHandler: OrpheusEventHandler {
    @Dependency(\.date) private var date
    private let partialJsonExtractor: PartialJsonExtractor

    public init(partialJsonExtractor: PartialJsonExtractor) {
        self.partialJsonExtractor = partialJsonExtractor
    }
    
    public func handle(
        event: OrpheusRealtimeEvent,
        session: OrpheusSession
    ) -> EventHandlerResult {
        guard case .messageToolCall(let messageId, let toolCallId, let toolCallName, let toolCallArguments) = event else {
            return .noOp
        }
        
        print("📡 handleMessageToolCall: messageId=\(messageId), toolCallId=\(toolCallId ?? "nil"), toolCallName=\(toolCallName ?? "nil")")
        
        var messagesList = session.messages
        
        // First, try to find message with the same messageId
        var existingMessage = messagesList.first { $0.id == messageId }
        var messageIndex: Int?
        
        // If not found and this is a simple_message tool call, try to find a recent assistant message
        // with generatedClips that should be merged with this reply text
        // This handles the case where SSE sends different messageIds but they're part of the same response
        // Key insight: if SSE events have the same messageId, they should be the same OrpheusMessage
        // But if they have different messageIds, we still merge simple_message into the message with clips
        if existingMessage == nil {
            let toolName = ToolName.fromValue(toolCallName)
            if toolName == nil || toolCallName == "simple_message" {
                // Look for the most recent assistant message with generatedClips
                // This is likely the message we should merge the reply text into
                // We merge based on: same response = same message, even if SSE sends different messageIds
                if let recentMessageWithClips = messagesList.last(where: { 
                    $0.role == .assistant && 
                    !$0.generatedClips.isEmpty
                }) {
                    existingMessage = recentMessageWithClips
                    messageIndex = messagesList.firstIndex { $0.id == recentMessageWithClips.id }
                    print("🔗 Merging simple_message (messageId: \(messageId)) into existing message with clips (id: \(recentMessageWithClips.id))")
                    // Use the existing message's ID so they're treated as the same message
                    // This ensures Chat.swift sees them as one message
                }
            }
        } else {
            messageIndex = messagesList.firstIndex { $0.id == messageId }
        }
        
        // Match Android: always accumulate arguments from existing tool call, even if different toolCallId
        // Android: val previousArguments = existingMessage?.accumulatingToolCall?.functionCall?.arguments ?: ""
        //          val accumulatedArguments = previousArguments + (event.toolCallArguments?.takeIf { it != "null" } ?: "")
        let previousArguments = existingMessage?.accumulatingToolCall?.functionCall.arguments ?? ""
        let toolCallArgs = toolCallArguments == "null" ? "" : (toolCallArguments ?? "")
        let accumulatedArguments = previousArguments + toolCallArgs
        
        // Extract content from accumulated arguments (only works for simple_message tool calls with "message" field)
        let newPartialText = extractPartialMessageContent(
            previousArguments: previousArguments,
            accumulatedArguments: accumulatedArguments
        )
        
        if let existing = existingMessage, let index = messageIndex {
            // Match Android: always update/replace the tool call (Android's createOrUpdateToolCall replaces it)
            let accumulatingToolCall = createOrUpdateToolCall(
                event: event,
                existingToolCall: existing.accumulatingToolCall,
                accumulatedArguments: accumulatedArguments
            )
            
            // When merging simple_message into a message with generatedClips, preserve the generatedClips
            // This ensures we don't lose the clips when updating with the reply text
            updateExistingMessage(
                messagesList: &messagesList,
                messageIndex: index,
                newPartialText: newPartialText,
                accumulatingToolCall: accumulatingToolCall,
                existingMessage: existing
            )
        } else {
            let assistantMessage = createNewAssistantMessage(
                event: event,
                messageId: messageId,
                newPartialText: newPartialText,
                accumulatedArguments: accumulatedArguments
            )
            messagesList.append(assistantMessage)
        }
        
        return .messagesUpdated(messagesList)
    }
    
    private func createOrUpdateToolCall(
        event: OrpheusRealtimeEvent,
        existingToolCall: ToolCall?,
        accumulatedArguments: String
    ) -> ToolCall {
        guard case .messageToolCall(_, let toolCallId, let toolCallName, _) = event else {
            fatalError("Invalid event type")
        }
        if let existing = existingToolCall {
            return ToolCall(
                id: existing.id,
                functionCall: FunctionCall(
                    name: existing.functionCall.name,
                    arguments: accumulatedArguments
                )
            )
        } else {
            return ToolCall(
                id: toolCallId,
                functionCall: FunctionCall(
                    name: toolCallName,
                    arguments: accumulatedArguments
                )
            )
        }
    }
    
    private func updateExistingMessage(
        messagesList: inout IdentifiedArrayOf<OrpheusMessage>,
        messageIndex: Int,
        newPartialText: String,
        accumulatingToolCall: ToolCall,
        existingMessage: OrpheusMessage
    ) {
        // Match Android: use extracted messageContent if not empty, otherwise keep existing content
        // Android: content = messageContent.ifEmpty { message.content }
        let finalContent = newPartialText.isEmpty ? existingMessage.content : newPartialText
        
        // When updating, preserve generatedClips if they exist (important when merging simple_message into message with clips)
        // The copy method preserves generatedClips by default, so we don't need to explicitly set it
        messagesList[messageIndex] = existingMessage.copy(
            content: finalContent,
            accumulatingToolCall: accumulatingToolCall
        )
    }
    
    private func createNewAssistantMessage(
        event: OrpheusRealtimeEvent,
        messageId: String,
        newPartialText: String,
        accumulatedArguments: String
    ) -> OrpheusMessage {
        guard case .messageToolCall(_, let toolCallId, let toolCallName, _) = event else {
            fatalError("Invalid event type")
        }
        
        let toolName = ToolName.fromValue(toolCallName)
        let contentType: MessageContentType
        if toolName == .writeLyrics {
            contentType = .toolCallContent
        } else {
            contentType = .chat
        }
        
        return OrpheusMessage(
            id: messageId,
            role: .assistant,
            content: newPartialText,
            timestamp: date.now.timeIntervalSince1970,
            contentType: contentType,
            accumulatingToolCall: ToolCall(
                id: toolCallId,
                functionCall: FunctionCall(
                    name: toolCallName,
                    arguments: accumulatedArguments
                )
            ),
            generatedClips: [],
            isStreaming: true
        )
    }
    
    private func extractPartialMessageContent(
        previousArguments: String,
        accumulatedArguments: String
    ) -> String {
        // Match Android: extract the full message content from accumulated arguments
        // Android: val messageContent = partialJsonExtractor.extractField(...) ?: ""
        // We use the full extracted message, not the difference
        return partialJsonExtractor.extractField(
            partialJson: accumulatedArguments,
            fieldName: JSON_FIELD_MESSAGE
        ) ?? ""
    }
}