"""
LiveKit Agents integration for automatic usage tracking with Paygent.
This hook intercepts LiveKit metrics and automatically sends usage data to Paygent.
Supports both the legacy VoicePipelineAgent and the newer AgentSession API.
"""

import logging
from typing import Any, Optional, Dict
from ..client import Client
from ..models import LLMConfig, STTConfig, TTSConfig, CustomerCreateOrGetRequest

class LiveKitPaygentHook:
    """
    Hook for LiveKit Agents that provides automatic usage tracking.
    Supports both legacy VoicePipelineAgent and the newer AgentSession/inference API.
    """

    def __init__(
        self, 
        paygent_client: Client, 
        agent_id: str, 
        customer_id: str, 
        session_id: str,
        indicator_name: str = "voice-session",
        customer_name: Optional[str] = None,
        llm_config: Optional[LLMConfig] = None,
        stt_config: Optional[STTConfig] = None,
        tts_config: Optional[TTSConfig] = None
    ):
        """
        Create a new LiveKitPaygentHook.
        
        Args:
            paygent_client: The Paygent client instance
            agent_id: Unique identifier for your agent (External ID)
            customer_id: Unique identifier for your customer
            session_id: The LiveKit session ID or a unique ID for this call
            indicator_name: The indicator used for billing (e.g., "voice-session")
            customer_name: Optional name for the customer to create/get
            llm_config: Optional manual configuration for LLM
            stt_config: Optional manual configuration for STT
            tts_config: Optional manual configuration for TTS
        """
        self.paygent = paygent_client
        self.agent_id = agent_id
        self.customer_id = customer_id
        self.customer_name = customer_name
        self.session_id = session_id
        self.indicator_name = indicator_name
        self.llm_config = llm_config
        self.stt_config = stt_config
        self.tts_config = tts_config
        self.logger = paygent_client.get_logger()
        self._initialized = False
        self._finalized = False
        self._start_time = None

    def attach(self, target: Any) -> None:
        """
        Attach this hook to a LiveKit Agent or AgentSession.
        
        Args:
            target: The LiveKit Agent, VoicePipelineAgent, or AgentSession instance
        """
        # Initialize voice session in Paygent backend
        try:
            # Create or get the customer in Paygent
            try:
                self.paygent.create_or_get_customer(
                    CustomerCreateOrGetRequest(
                        name=self.customer_name or self.customer_id,
                        external_id=self.customer_id
                    )
                )
            except Exception as e:
                self.logger.warning(f"Failed to create/get customer in Paygent: {e}")

            self.paygent.initialize_voice_session(
                self.session_id, 
                self.agent_id, 
                self.customer_id
            )
            self._initialized = True
        except Exception as e:
            self.logger.error(f"Failed to initialize Paygent voice session: {e}")

        # Automatically track duration and end session when target stops/shutdowns
        # Note: _start_time will be set when the first metric is received or when session starts
        
        # For AgentSession (v1.4+), hook into individual inference components
        if hasattr(target, 'llm') and hasattr(target, 'stt') and hasattr(target, 'tts'):
            
            # Hook into LLM metrics
            if target.llm:
                @target.llm.on("metrics_collected")
                def _on_llm_metrics(metrics: Any):
                    self._handle_llm_metrics(metrics)
            
            # Hook into STT metrics
            if target.stt:
                @target.stt.on("metrics_collected")
                def _on_stt_metrics(metrics: Any):
                    self._handle_stt_metrics(metrics)
            
            # Hook into TTS metrics
            if target.tts:
                @target.tts.on("metrics_collected")
                def _on_tts_metrics(metrics: Any):
                    self._handle_tts_metrics(metrics)
        
        else:
            # Legacy VoicePipelineAgent or older API
            @target.on("metrics_collected")
            def _on_metrics_collected(metrics: Any):
                self._handle_metrics(metrics)
        
        
        # Listen to multiple session end events to ensure indicator is triggered
        # Different LiveKit versions may use different event names
        
        @target.on("shutdown")
        def _on_shutdown():
            self.end_session()
        
        # Also listen to 'close' event (emitted by AgentSession)
        try:
            @target.on("close")
            def _on_close(*args, **kwargs):
                self.end_session()
        except Exception as e:
            pass

    def _handle_metrics(self, metrics: Any) -> None:
        """
        Internal handler for LiveKit metrics.
        Maps LiveKit metric objects to Paygent tracking calls.
        """
        try:
            # Check for LLM metrics
            if hasattr(metrics, 'prompt_tokens') and hasattr(metrics, 'completion_tokens'):
                self._handle_llm_metrics(metrics)
            
            # Check for STT metrics
            elif hasattr(metrics, 'audio_duration') and not hasattr(metrics, 'playback_duration'):
                self._handle_stt_metrics(metrics)
                
            # Check for TTS metrics
            elif hasattr(metrics, 'audio_duration') and hasattr(metrics, 'playback_duration'):
                self._handle_tts_metrics(metrics)
                
        except Exception as e:
            self.logger.error(f"Error handling LiveKit metrics in Paygent hook: {e}")

    def _handle_llm_metrics(self, metrics: Any) -> None:
        """Process and send LLM usage."""
        # Set start time on first metrics event if not already set
        if not self._start_time:
            import time
            self._start_time = time.time()
        
        if self.llm_config:
            provider = self.llm_config.provider
            model = self.llm_config.model
            plan = self.llm_config.plan
        else:
            # Extract fields robustly if no manual config
            label = getattr(metrics, 'label', None)
            model_attr = getattr(metrics, 'model', None)
            
            model_str = "unknown/unknown"
            
            # If label or model is an object with a .model property (like the LLM class itself)
            if hasattr(label, 'model') and isinstance(label.model, str):
                model_str = label.model
            elif hasattr(model_attr, 'model') and isinstance(model_attr.model, str):
                model_str = model_attr.model
            elif isinstance(label, str):
                model_str = label
            elif isinstance(model_attr, str):
                model_str = model_attr
                
            # Parse provider and model
            if '/' in str(model_str):
                parts = str(model_str).split('/', 1)
                provider = parts[0]
                model = parts[1]
            else:
                provider = 'unknown'
                model = str(model_str)
            plan = "default"
        
        prompt_tokens = getattr(metrics, 'prompt_tokens', 0) or getattr(metrics, 'input_tokens', 0)
        completion_tokens = getattr(metrics, 'completion_tokens', 0) or getattr(metrics, 'output_tokens', 0)
        
        llm_data = {
            "serviceProvider": provider,
            "model": model,
            "promptTokens": prompt_tokens,
            "completionTokens": completion_tokens,
            "plan": plan
        }
        
        self.paygent.track_voice_llm(self.session_id, llm_data)

    def _handle_stt_metrics(self, metrics: Any) -> None:
        """Process and send STT usage."""
        # Set start time on first metrics event if not already set
        if not self._start_time:
            import time
            self._start_time = time.time()
        
        if self.stt_config:
            provider = self.stt_config.provider
            model = self.stt_config.model
            plan = self.stt_config.plan
        else:
            label = getattr(metrics, 'label', None)
            model_attr = getattr(metrics, 'model', None)
            
            model_str = "unknown/unknown"
            if hasattr(label, 'model') and isinstance(label.model, str):
                model_str = label.model
            elif hasattr(model_attr, 'model') and isinstance(model_attr.model, str):
                model_str = model_attr.model
            elif isinstance(label, str):
                model_str = label
            elif isinstance(model_attr, str):
                model_str = model_attr

            if '/' in str(model_str):
                parts = str(model_str).split('/', 1)
                provider = parts[0]
                model = parts[1]
            else:
                provider = 'unknown'
                model = str(model_str)
            plan = "default"
            
        # LiveKit audio_duration is in seconds, Paygent expects minutes
        audio_duration = getattr(metrics, 'audio_duration', 0) or getattr(metrics, 'duration', 0)
        audio_minutes = audio_duration / 60.0
        
        stt_data = {
            "serviceProvider": provider,
            "model": model,
            "audioDuration": audio_minutes,
            "plan": plan
        }
        
        self.paygent.track_voice_stt(self.session_id, stt_data)

    def _handle_tts_metrics(self, metrics: Any) -> None:
        """Process and send TTS usage."""
        # Set start time on first metrics event if not already set
        if not self._start_time:
            import time
            self._start_time = time.time()
        
        if self.tts_config:
            provider = self.tts_config.provider
            model = self.tts_config.model
            plan = self.tts_config.plan
        else:
            label = getattr(metrics, 'label', None)
            model_attr = getattr(metrics, 'model', None)
            
            model_str = "unknown/unknown"
            if hasattr(label, 'model') and isinstance(label.model, str):
                model_str = label.model
            elif hasattr(model_attr, 'model') and isinstance(model_attr.model, str):
                model_str = model_attr.model
            elif isinstance(label, str):
                model_str = label
            elif isinstance(model_attr, str):
                model_str = model_attr

            if '/' in str(model_str):
                parts = str(model_str).split('/', 1)
                provider = parts[0]
                model = parts[1]
            else:
                provider = 'unknown'
                model = str(model_str)
            plan = "default"
        
        # Try multiple possible field names for character count
        char_count = (
            getattr(metrics, 'character_count', None) or
            getattr(metrics, 'characters_count', None) or
            getattr(metrics, 'text_length', None) or
            getattr(metrics, 'characters', None) or
            0
        )
        
        # Skip if character count is 0 or None
        if not char_count or char_count == 0:
            self.logger.warning(f"Skipping TTS tracking - character count is {char_count}")
            return
        
        tts_data = {
            "serviceProvider": provider,
            "model": model,
            "characterCount": char_count,
            "plan": plan
        }
        
        try:
            self.paygent.track_voice_tts(self.session_id, tts_data)
        except Exception as e:
            self.logger.error(f"Failed to track TTS: {e}")

    def end_session(self) -> None:
        """
        Finalize the session and report the indicator for billing.
        This method is idempotent and can be safely called multiple times.
        """
        if self._finalized:
            return

        if not self._initialized or not self._start_time:
            self.logger.warning(f"⚠️ Skipping end_session - not initialized or no start time")
            # Mark as finalized anyway if it's never going to have a start time
            if not self._initialized:
                self._finalized = True
            return

        import time
        duration_seconds = time.time() - self._start_time
        duration_minutes = duration_seconds / 60.0
        
        try:
            self.paygent.set_voice_indicator(
                self.session_id,
                self.indicator_name,
                duration_minutes
            )
            self.logger.info(f"Paygent session finalized with indicator '{self.indicator_name}' (Duration: {duration_minutes:.2f}m)")
            self._finalized = True
            self._initialized = False 
        except Exception as e:
            self.logger.error(f"❌ Failed to finalize Paygent session: {e}")
