package com.saas.voip.controller;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.genai.Client;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.GenerateContentResponse;
import com.google.genai.types.Part;
import com.saas.shared.dto.AiCostCalculationResult;
import com.saas.shared.service.PhoneNumberBoundaryService;
import com.saas.shared.service.ai.AiCostTrackingService;
import com.saas.voip.service.EmmaPromptService;
import com.saas.voip.service.ai.GeminiCostCalculator;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Controller;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Gemini ConversationRelay Controller for Twilio
 * 
 * Uses Twilio's ConversationRelay feature with Gemini Text API (not Live API).
 * This is the SIMPLE approach - no audio conversion, just text exchange.
 * 
 * Architecture:
 * - Twilio handles TTS/STT via ConversationRelay
 * - We exchange JSON text messages via WebSocket
 * - Gemini Text API processes prompts (not audio)
 * - Much simpler than GeminiRealtimeService!
 * 
 * Based on Python example provided by user.
 */
@Controller
@Slf4j
@RequiredArgsConstructor
public class GeminiConversationRelayController extends TextWebSocketHandler {

    private final ObjectMapper objectMapper;
    private final EmmaPromptService emmaPromptService;
    private final GeminiCostCalculator geminiCostCalculator;
    private final AiCostTrackingService aiCostTrackingService;
    private final PhoneNumberBoundaryService phoneNumberBoundaryService;
    
    @Value("${gemini.api.key}")
    private String geminiApiKey;
    
    @Value("${gemini.model:gemini-2.0-flash}")
    private String geminiModel;
    
    // Store active chat sessions per call (Client + conversation history)
    private final Map<String, Client> geminiSessions = new ConcurrentHashMap<>();
    private final Map<String, List<Content>> conversationHistories = new ConcurrentHashMap<>();
    
    // Track usage and metadata per session for cost calculation
    private final Map<String, SessionUsage> sessionUsageMap = new ConcurrentHashMap<>();
    
    @Data
    private static class SessionUsage {
        private String callSid;
        private String tenantId;
        private LocalDateTime startTime;
        private LocalDateTime endTime;
        private long textInputTokens = 0;
        private long textOutputTokens = 0;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("🔗 [Gemini ConversationRelay] WebSocket connection established: {}", session.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        try {
            String payload = message.getPayload();
            log.info("📨 [Gemini ConversationRelay] Raw message received: {}", payload);
            
            JsonNode jsonMessage = objectMapper.readTree(payload);
            String messageType = jsonMessage.get("type").asText();
            
            log.info("📩 [Gemini ConversationRelay] Message type: {}", messageType);

            switch (messageType) {
                case "setup":
                    handleSetup(session, jsonMessage);
                    break;
                case "prompt":
                    handlePrompt(session, jsonMessage);
                    break;
                case "interrupt":
                    handleInterrupt(session, jsonMessage);
                    break;
                default:
                    log.warn("⚠️ [Gemini ConversationRelay] Unknown message type: {}", messageType);
                    log.warn("⚠️ [Gemini ConversationRelay] Full message: {}", payload);
            }
            
        } catch (Exception e) {
            log.error("❌ [Gemini ConversationRelay] Error handling message: {}", message.getPayload(), e);
        }
    }

    /**
     * Handle setup message from Twilio ConversationRelay
     */
    private void handleSetup(WebSocketSession session, JsonNode message) {
        String callSid = message.get("callSid").asText();
        log.info("🎤 [Gemini ConversationRelay] Setup for call: {}", callSid);
        
        try {
            // Initialize Gemini client
            Client geminiClient = Client.builder()
                    .apiKey(geminiApiKey)
                    .build();
            
            // Initialize conversation history with system prompt
            List<Content> history = new ArrayList<>();
            
            // Store session and history
            geminiSessions.put(session.getId(), geminiClient);
            conversationHistories.put(session.getId(), history);
            
            // Initialize usage tracking
            SessionUsage usage = new SessionUsage();
            usage.setCallSid(callSid);
            usage.setStartTime(LocalDateTime.now());
            
            // Resolve tenantId from phone number (available in setup message)
            String toNumber = message.has("to") ? message.get("to").asText() : null;
            String tenantId = resolveTenantIdFromPhoneNumber(toNumber);
            usage.setTenantId(tenantId);
            
            sessionUsageMap.put(session.getId(), usage);
            
            log.info("✅ [Gemini ConversationRelay] Gemini client initialized for call: {}", callSid);
            log.info("📋 [Gemini ConversationRelay] Conversation history initialized with system prompt");
            log.info("💰 [Gemini ConversationRelay] Usage tracking started (tenantId: {})", tenantId);
            
        } catch (Exception e) {
            log.error("❌ [Gemini ConversationRelay] Error initializing Gemini client", e);
        }
    }

    /**
     * Handle prompt (user speech transcribed by Twilio)
     */
    private void handlePrompt(WebSocketSession session, JsonNode message) throws Exception {
        String userPrompt = message.get("voicePrompt").asText();
        log.info("💬 [Gemini ConversationRelay] Processing prompt: {}", userPrompt);
        
        Client geminiClient = geminiSessions.get(session.getId());
        List<Content> history = conversationHistories.get(session.getId());
        
        if (geminiClient == null || history == null) {
            log.error("❌ [Gemini ConversationRelay] No Gemini client or history for session: {}", session.getId());
            return;
        }
        
        try {
            // Add user message to history
            history.add(Content.builder()
                    .role("user")
                    .parts(List.of(Part.fromText(userPrompt)))
                    .build());
            
            log.info("📝 [Gemini ConversationRelay] History size: {} messages", history.size());
            
            // Create config with Emma system instruction (prompt + RAG)
            GenerateContentConfig config = GenerateContentConfig.builder()
                    .systemInstruction(Content.fromParts(Part.fromText(emmaPromptService.getFullSystemPrompt())))
                    .build();
            
            // Get response from Gemini with full conversation history
            GenerateContentResponse response = geminiClient.models.generateContent(
                    geminiModel,
                    history,
                    config
            );
            
            String responseText = response.text();
            log.info("🤖 [Gemini ConversationRelay] Gemini response: {}", responseText);
            
            // Track token usage for cost calculation
            SessionUsage usage = sessionUsageMap.get(session.getId());
            if (usage != null) {
                if (response.usageMetadata() != null && response.usageMetadata().isPresent()) {
                    var metadata = response.usageMetadata().get();
                    long inputTokens = metadata.promptTokenCount().orElse(0);
                    long outputTokens = metadata.candidatesTokenCount().orElse(0);
                    
                    usage.setTextInputTokens(usage.getTextInputTokens() + inputTokens);
                    usage.setTextOutputTokens(usage.getTextOutputTokens() + outputTokens);
                    
                    log.info("💰 [Gemini ConversationRelay] Token usage - Input: {}, Output: {} (Total: input={}, output={})", 
                             inputTokens, outputTokens, usage.getTextInputTokens(), usage.getTextOutputTokens());
                } else {
                    log.warn("⚠️ [Gemini ConversationRelay] No usage metadata in response! Tokens won't be tracked.");
                    log.warn("⚠️ [Gemini ConversationRelay] Response text length: {} chars", responseText != null ? responseText.length() : 0);
                    log.warn("⚠️ [Gemini ConversationRelay] This will result in ZERO cost tracking!");
                }
            } else {
                log.error("❌ [Gemini ConversationRelay] No usage tracking session found for session: {}", session.getId());
            }
            
            // Add AI response to history
            history.add(Content.builder()
                    .role("model")
                    .parts(List.of(Part.fromText(responseText)))
                    .build());
            
            // Send response back to Twilio
            // Twilio will convert text to speech via ConversationRelay
            String jsonResponse = objectMapper.writeValueAsString(Map.of(
                    "type", "text",
                    "token", responseText,
                    "last", true  // Indicate full and final message
            ));
            
            session.sendMessage(new TextMessage(jsonResponse));
            log.info("📤 [Gemini ConversationRelay] Response sent to Twilio (history: {} messages)", history.size());
            
        } catch (Exception e) {
            log.error("❌ [Gemini ConversationRelay] Error getting Gemini response", e);
        }
    }

    /**
     * Handle interrupt (user interrupted AI)
     */
    private void handleInterrupt(WebSocketSession session, JsonNode message) {
        String callSid = message.has("callSid") ? message.get("callSid").asText() : "unknown";
        log.info("⏸️ [Gemini ConversationRelay] Handling interruption for call: {}", callSid);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        List<Content> history = conversationHistories.get(session.getId());
        int messageCount = (history != null) ? history.size() : 0;
        
        log.info("🔌 [Gemini ConversationRelay] Connection closed: {} - Status: {} - Total messages: {}", 
                session.getId(), status, messageCount);
        
        // Save AI costs before cleanup
        try {
            saveAiCostsForSession(session.getId());
        } catch (Exception e) {
            log.error("❌ [Gemini ConversationRelay] Error saving AI costs", e);
        }
        
        geminiSessions.remove(session.getId());
        conversationHistories.remove(session.getId());
        sessionUsageMap.remove(session.getId());
    }
    
    /**
     * Resolve tenant ID from phone number using boundary service
     * Maintains clean architecture by not directly accessing admin repository
     */
    private String resolveTenantIdFromPhoneNumber(String phoneNumber) {
        String tenantId = phoneNumberBoundaryService.resolveTenantIdFromPhoneNumber(phoneNumber);
        if (tenantId == null) {
            log.warn("⚠️ [Gemini ConversationRelay] No tenant found for phone {}, costs will be saved without tenant attribution", phoneNumber);
        } else {
            log.info("✅ [Gemini ConversationRelay] Tenant resolved for phone {}: {}", phoneNumber, tenantId);
        }
        return tenantId;
    }
    
    /**
     * Calculate and save AI costs to admin database
     */
    private void saveAiCostsForSession(String sessionId) {
        SessionUsage usage = sessionUsageMap.get(sessionId);
        if (usage == null) {
            log.warn("⚠️ [Gemini ConversationRelay] No usage data found for session: {}", sessionId);
            return;
        }
        
        usage.setEndTime(LocalDateTime.now());
        
        // Check if any tokens were used
        long totalTokens = usage.getTextInputTokens() + usage.getTextOutputTokens();
        if (totalTokens == 0) {
            log.info("💰 [Gemini ConversationRelay] No token usage detected for session: {}, skipping cost tracking", sessionId);
            return;
        }
        
        // Prepare usage metrics for calculator
        Map<String, Object> usageMetrics = new HashMap<>();
        usageMetrics.put("text_input_tokens", usage.getTextInputTokens());
        usageMetrics.put("text_output_tokens", usage.getTextOutputTokens());
        usageMetrics.put("model", geminiModel);
        
        // Calculate cost using GeminiCostCalculator
        AiCostCalculationResult costResult = geminiCostCalculator.calculateCost(usageMetrics);
        
        if (costResult != null && costResult.isValid()) {
            log.info("💰 [Gemini ConversationRelay] Cost calculated - ${} USD for call {}", 
                     costResult.getCost(), usage.getCallSid());
            
            // Save to admin database
            aiCostTrackingService.saveAiCost(
                usage.getCallSid(),
                usage.getTenantId(),
                costResult,
                usage.getStartTime(),
                usage.getEndTime(),
                null, // fromNumber
                null  // toNumber
            );
            
            log.info("✅ [Gemini ConversationRelay] AI costs saved to admin database for call {}", usage.getCallSid());
        } else {
            log.warn("⚠️ [Gemini ConversationRelay] Invalid cost calculation result for session: {}", sessionId);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("❌ [Gemini ConversationRelay] Transport error for session: {}", session.getId(), exception);
        
        // Try to save costs even on error
        try {
            saveAiCostsForSession(session.getId());
        } catch (Exception e) {
            log.error("❌ [Gemini ConversationRelay] Error saving AI costs during transport error", e);
        }
        
        geminiSessions.remove(session.getId());
        conversationHistories.remove(session.getId());
        sessionUsageMap.remove(session.getId());
    }
}