from dataclasses import is_dataclass
import json
import logging
import os
from datetime import datetime
from typing import Dict, List, Any, Optional
import asyncio

import aiohttp
from dotenv import load_dotenv
from livekit.agents import (
    AgentSession,
    JobContext,
    JobProcess,
    WorkerOptions,
    MetricsCollectedEvent,
    RoomInputOptions,
    cli,
    metrics,
    room_io,
)
from livekit.plugins import silero
from livekit.plugins import noise_cancellation
from livekit.plugins.turn_detector.multilingual import MultilingualModel
from livekit.plugins import elevenlabs
from livekit.plugins import openai
from livekit.plugins import deepgram
from livekit.plugins import bey
from livekit import api
from livekit import rtc

# Import our agent factory and related classes
from agents import AgentFactory

logger = logging.getLogger("agent")

load_dotenv(".env.local")

async def record_room_to_s3(room_name: str):  
    # Initialize LiveKit API client  
    lkapi = api.LiveKitAPI()  
      
    # Configure S3 upload  
    s3_upload = api.S3Upload(  
        access_key=os.environ.get('AWS_ACCESS_KEY_ID'),  
        secret=os.environ.get('AWS_SECRET_ACCESS_KEY'),  
        bucket=os.environ.get('S3_RECORDING_BUCKET'),  
        region=os.environ.get('AWS_DEFAULT_REGION', 'ap-south-1')  
    )  
      
    # Start room composite recording  
    egress_request = api.RoomCompositeEgressRequest(  
        room_name=room_name,  
        layout='speaker',  # or 'grid' layout  
        preset=api.EncodingOptionsPreset.H264_720P_30,  
        audio_only=True,  
        file_outputs=[api.EncodedFileOutput(  
            file_type=api.EncodedFileType.MP3,  
            filepath=f'recordings/{room_name}/{room_name}.mp3',  
            s3=s3_upload  
        )]  
    )  
      
    # Start the egress  
    egress_info = await lkapi.egress.start_room_composite_egress(egress_request)  

    print(f'Egress started: {egress_info.egress_id}')  
      
    await lkapi.aclose()  
    return f"recordings/{room_name}/{room_name}.mp3" if egress_info else ""

# agent_config
def build_stt(provider: str, model: str, **kwargs):  
    if provider == "deepgram":  
        return deepgram.STTv2(model=model, **kwargs)  
    elif provider == "openai":  
        return openai.STT(model=model, **kwargs)  
    # Add more providers
    else:
        raise ValueError(f"Unsupported STT provider: {provider}")  
      
def build_llm(provider: str, model: str, **kwargs):  
    if provider == "openai":  
        return openai.LLM(model=model, **kwargs)  
    # Add more providers
    else:
        raise ValueError(f"Unsupported LLM provider: {provider}")  
      
def build_tts(provider: str, voice_id: str, model: str, **kwargs):  
    if provider == "elevenlabs":  
        return elevenlabs.TTS(voice_id=voice_id, model=model, **kwargs)  
    # Add more providers
    else:
        raise ValueError(f"Unsupported TTS provider: {provider}")

def prewarm(proc: JobProcess):
    proc.userdata["vad"] = silero.VAD.load()

def to_json_safe(value):
    """Recursively convert an object to JSON-serializable format."""
    from dataclasses import is_dataclass, asdict

    if value is None:
        return None

    # Already JSON safe
    if isinstance(value, (str, int, float, bool)):
        return value

    # List/tuple → convert each element
    if isinstance(value, (list, tuple)):
        return [to_json_safe(v) for v in value]

    # Dict → convert values
    if isinstance(value, dict):
        return {k: to_json_safe(v) for k, v in value.items()}

    # Dataclass (e.g., Metadata)
    if is_dataclass(value):
        return to_json_safe(asdict(value))

    # Objects with __dict__
    if hasattr(value, "__dict__"):
        return to_json_safe({
            k: v for k, v in value.__dict__.items() if not k.startswith("_")
        })

    # Fallback
    return str(value)

def serialize_metrics(metrics_obj: Any) -> Optional[Dict[str, Any]]:
    if metrics_obj is None:
        return None

    try:
        result = {}

        # copy known attributes
        for attr in [
            'type', 'speech_id', 'duration', 'tokens', 'input_tokens',
            'output_tokens', 'model'
        ]:
            if hasattr(metrics_obj, attr):
                result[attr] = to_json_safe(getattr(metrics_obj, attr))

        # add rest of public __dict__
        if hasattr(metrics_obj, "__dict__"):
            for k, v in metrics_obj.__dict__.items():
                if not k.startswith("_"):
                    result[k] = to_json_safe(v)

        return result

    except Exception as e:
        # logger.warning(f"Failed to serialize metrics: {e}")
        return None


def serialize_metrics_list(metrics_list: List[Any]) -> List[Dict[str, Any]]:
    """Convert list of metrics objects to JSON-serializable list"""
    if not metrics_list:
        return []
    return [m for m in [serialize_metrics(metric) for metric in metrics_list] if m is not None]

async def create_conversation_from_backend(agent_id: str, recordingPath: str) -> str | None:
    """Call backend API to create a new conversation and return its ID."""
    backend_url = os.environ.get("AGENT_API_BASE_URL")
    if not backend_url:
        # logger.warning("AGENT_API_BASE_URL environment variable not set, skipping conversation upload")
        return None
    payload = {"agentId": agent_id, "recordingPath": recordingPath or ""}
    # logger.info(f"Creating conversation for agent: {agent_id}")
    try:
        async with aiohttp.ClientSession() as http_session:
            async with http_session.post(
                f"{backend_url}/v1/conversation",
                json=payload,
                timeout=aiohttp.ClientTimeout(total=10)
            ) as resp:

                response_text = await resp.text()

                if resp.status == 200:
                    try:
                        data = json.loads(response_text)
                        conversation_id = data.get("data", {}).get("conversation_id")
                        # logger.debug(f"Backend response data: {data}")
                        if conversation_id:
                            # logger.info(f"Successfully created conversation {conversation_id} for agent: {agent_id}")
                            return conversation_id
                        else:
                            # logger.warning(f"No conversationId in backend response: {data}")
                            return None

                    except json.JSONDecodeError:
                        # logger.error("Backend returned invalid JSON")
                        return None

                else:
                    # logger.warning(f"Backend returned status {resp.status}: {response_text}")
                    return None

    except Exception as e:
        # logger.error(f"Failed to send conversation to backend: {e}", exc_info=True)
        return None
     
       
async def entrypoint(ctx: JobContext):  
    await ctx.connect()  

    participant = await ctx.wait_for_participant()
    
    # Get the agent ID from participant attributes
    # The agent ID is used to fetch all agent configuration dynamically from the API
    agent_id = participant.attributes.get("agent")
    # Default to the fallback agent when no attribute is provided
    if not agent_id:
        agent_id = "693a96fd276eb7171e083efb" # Dbles Inbound Agent
    user_id = "69299a964990c460f4188cf3"
    # user_id = participant.attributes.get("userId")
    
    if not agent_id:
        # logger.error("No agent ID provided in participant attributes")
        raise ValueError("Agent ID is required. Please provide 'agent' attribute in participant.")
    
    # logger.info(f"Creating dynamic agent with agent_id: {agent_id}")

    # Add any other context you want in all log entries here
    ctx.log_context_fields = {
        "room": ctx.room.name,
        "agent_type": "dynamic_rag_agent",
        "agent_id": agent_id,
    }
    egress_info = ""

  # Record the room to S3
    try:
        if (os.getenv('S3_RECORDING_BUCKET') is not None):
            # logger.info(f"Recording room to S3: {ctx.room.name}")
            egress_info = await record_room_to_s3(ctx.room.name)
        else:
            # logger.info(f"S3 recording bucket is not configured, skipping recording")
            pass
    except Exception as e:
        # logger.error(f"Failed to record room to S3: {e}")
        pass
        
    conversation_id = await create_conversation_from_backend(agent_id, egress_info)
    try:
        # Create dynamic agent - all configuration will be fetched from API
        agent = await AgentFactory.create_agent(
            "dynamic_rag_agent", agent_id=agent_id
        )
        config = agent.config
        config = config["formatted_config"]
        
        # logger.info(f"Created agent: {config['agent_name']}")
        
    except ValueError as e:
        # logger.error(f"Failed to create agent: {e}")
        raise  # Re-raise to prevent continuing with invalid agent

    # Build session components using the agent's configuration
    session = AgentSession(
        stt=config["stt"],
        llm=config["llm"],
        tts=config["tts"],
        turn_detection=MultilingualModel(),
        vad=ctx.proc.userdata["vad"],
        preemptive_generation=True,
    )

    # Metrics collection, to measure pipeline performance
    # For more information, see https://docs.livekit.io/agents/build/metrics/
    usage_collector = metrics.UsageCollector()
    last_eou_metrics: Optional[metrics.EOUMetrics] = None
    conversation_metrics: Dict[str, Dict[str, Any]] = {}  # Maps speech_id to collected metrics
    pending_stt_metrics: List[Any] = []  # Store STT metrics until we get EOUMetrics with speech_id
    conversation_messages: List[Dict[str, Any]] = []  # Store all conversation messages with usage data
    flow_start_times: Dict[str, float] = {}  # Maps speech_id to flow start time (perf_counter)

    @session.on("metrics_collected")
    def _on_metrics_collected(ev: MetricsCollectedEvent):
        nonlocal last_eou_metrics

        # Handle STT metrics (no speech_id)
        if ev.metrics.type == "stt_metrics":
            import time
            # Store temporarily until we get EOUMetrics
            pending_stt_metrics.append(ev.metrics)
            stt_metrics = ev.metrics
            audio_duration = getattr(stt_metrics, 'audio_duration', 0)
            metrics.log_metrics(ev.metrics)
            usage_collector.collect(ev.metrics)
            return

        # Get speech_id from the metric if available
        speech_id = getattr(ev.metrics, 'speech_id', None)

        if speech_id:
            if speech_id not in conversation_metrics:
                conversation_metrics[speech_id] = {
                    'llm': None,
                    'tts': [],
                    'stt': [],
                    'eou': None
                }

            # Store metrics by type
            if ev.metrics.type == "llm_metrics":
                import time
                from datetime import datetime
                llm_metrics = ev.metrics
                conversation_metrics[speech_id]['llm'] = llm_metrics
                
                # Log LLM processing time in seconds
                model_name = llm_metrics.metadata.model_name if hasattr(llm_metrics.metadata, 'model_name') else 'unknown'
            elif ev.metrics.type == "tts_metrics":
                import time
                tts_metrics = ev.metrics
                conversation_metrics[speech_id]['tts'].append(tts_metrics)
            elif ev.metrics.type == "eou_metrics":
                import time
                eou_metrics = ev.metrics
                # Associate pending STT metrics with this speech_id
                conversation_metrics[speech_id]['stt'] = pending_stt_metrics.copy()
                pending_stt_metrics.clear()
                conversation_metrics[speech_id]['eou'] = ev.metrics
                last_eou_metrics = ev.metrics

        metrics.log_metrics(ev.metrics)
        usage_collector.collect(ev.metrics)

    @session.on("conversation_item_added")
    def _on_conversation_item_added(ev):
        asyncio.create_task(handle_conversation_item(ev))

    async def handle_conversation_item(ev):
        import time
        start_time = time.perf_counter()
        nonlocal last_eou_metrics
        item = ev.item
        timestamp = datetime.utcnow().isoformat()

        if item.role == "user":
            # User message added - this typically triggers an LLM request
            user_content = item.content[0] if isinstance(item.content, list) else str(item.content[0])
            
            # Track flow start time
            flow_start_perf = time.perf_counter()
            flow_start_timestamp = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            
            if last_eou_metrics:
                speech_id = last_eou_metrics.speech_id
                flow_start_times[speech_id] = flow_start_perf
                logger.info(
                    f"⏱️ [FLOW START] speech_id={speech_id}, "
                    f"timestamp={flow_start_timestamp}, "
                    f"user_message=\"{user_content[:100]}\""
                )
                metrics_data = conversation_metrics.get(speech_id, {})
            else:
                # Fallback if no speech_id available yet
                speech_id = f"unknown_{int(time.time())}"
                flow_start_times[speech_id] = flow_start_perf
                logger.info(
                    f"⏱️ [FLOW START] speech_id={speech_id}, "
                    f"timestamp={flow_start_timestamp}, "
                    f"user_message=\"{user_content[:100]}\""
                )
                metrics_data = {}
            
            message_data = {
                "conversationId": conversation_id,
                "content": item.content[0],
                "role": item.role,
                "timeStamp": timestamp,
                "agentId": agent_id,
            }

            if metrics_data:
                message_data["usage"] = {
                    "stt": serialize_metrics_list(metrics_data.get("stt", [])),
                    "eou": serialize_metrics(metrics_data.get("eou")),
                }

            conversation_messages.append(message_data)
            # Run in background without blocking conversation flow
            asyncio.create_task(send_conversation_to_backend(message_data))

        elif item.role == "assistant" and session.current_speech:
            speech_id = session.current_speech.id
            metrics_data = conversation_metrics.get(speech_id, {})

            # LLM metrics
            llm_raw = metrics_data.get("llm")
            llm_filtered = None
            if llm_raw:
                llm_filtered = {
                    "duration": llm_raw.duration,
                    "timestamp": llm_raw.timestamp,
                    "ttft": llm_raw.ttft,
                    "completion_tokens": llm_raw.completion_tokens,
                    "prompt_tokens": llm_raw.prompt_tokens,
                    "prompt_cached_tokens": llm_raw.prompt_cached_tokens,
                    "total_tokens": llm_raw.total_tokens,
                    "tokens_per_second": llm_raw.tokens_per_second,
                    "model_name": llm_raw.metadata.model_name,
                    "model_provider": llm_raw.metadata.model_provider,
                }

            # TTS metrics
            tts_filtered = [
                {
                    "duration": t.duration,
                    "ttfb": t.ttfb,
                    "audio_duration": t.audio_duration,
                    "characters_count": t.characters_count,
                    "model_name": t.metadata.model_name,
                    "model_provider": t.metadata.model_provider,
                }
                for t in metrics_data.get("tts", [])
            ]

            usage_data = {}

            if llm_filtered:
                usage_data["llm"] = llm_filtered

            if tts_filtered:
                usage_data["tts"] = tts_filtered
            
            message_data = {
                "agentId": agent_id,
                "conversationId": conversation_id,
                "content": item.content[0],
                "role": item.role,
                "timeStamp": timestamp,
                "usage": usage_data
            }

            conversation_messages.append(message_data)
            # Run in background without blocking conversation flow
            asyncio.create_task(send_conversation_to_backend(message_data))
            
            # Log flow end with total elapsed time
            flow_end_perf = time.perf_counter()
            flow_end_timestamp = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
            assistant_content = item.content[0] if isinstance(item.content, list) else str(item.content[0])
            
            if speech_id in flow_start_times:
                total_flow_time = flow_end_perf - flow_start_times[speech_id]
                logger.info(
                    f"⏱️ [FLOW END] speech_id={speech_id}, "
                    f"timestamp={flow_end_timestamp}, "
                    f"total_flow_time={total_flow_time:.3f}s, "
                    f"assistant_response=\"{assistant_content[:100]}\""
                )
                # Clean up to avoid memory leak
                del flow_start_times[speech_id]
            else:
                # Fallback if flow start wasn't tracked (shouldn't happen normally)
                logger.info(
                    f"⏱️ [FLOW END] speech_id={speech_id}, "
                    f"timestamp={flow_end_timestamp}, "
                    f"total_flow_time=unknown (start not tracked), "
                    f"assistant_response=\"{assistant_content[:100]}\""
                )

    async def send_conversation_to_backend(message_data):
        """Send conversation data to backend API in fire-and-forget manner"""
        backend_url = os.environ.get("AGENT_API_BASE_URL")
        if not backend_url:
            # logger.warning("AGENT_API_BASE_URL environment variable not set, skipping conversation upload")
            return
        # logger.info(f"Sending conversation data to backend: {message_data}")
        try:
            async with aiohttp.ClientSession() as http_session:
                async with http_session.post(
                    f"{backend_url}/v1/conversation/chat",
                    json=message_data,
                    timeout=aiohttp.ClientTimeout(total=10)
                ) as resp:
                    if resp.status == 200:
                        # logger.info(f"Successfully sent conversation data to backend for agent: {agent_id}")
                        pass
                    else:
                        # logger.warning(f"Backend API returned status {resp.status}: {await resp.text()}")
                        pass
        except Exception as e:
            # logger.error(f"Failed to send conversation to backend: {e}", exc_info=True)
            pass

    async def send_end_call_to_backend():
        """Send end-of-call usage summary to backend /v1/conversation/end"""

        backend_url = os.environ.get("AGENT_API_BASE_URL")
        if not backend_url:
            # logger.warning("AGENT_API_BASE_URL environment variable not set, skipping end-call upload")
            return

        # ---- FETCH USAGE SUMMARY ----
        try:
            summary = usage_collector.get_summary()
            # logger.info(f"USAGE END CALL SUMMARY: {summary}")
            if not summary:
                # logger.warning("usage_collector.get_summary() returned no data")
                return

            usage_payload = {
                "llm": {
                    "completionTokens": summary.llm_completion_tokens,
                    "promptTokens": summary.llm_prompt_tokens,
                    "promptCachedTokens": summary.llm_prompt_cached_tokens,
                    "totalTokens": (
                        summary.llm_prompt_tokens +
                        summary.llm_prompt_cached_tokens +
                        summary.llm_completion_tokens
                    ),
                },
                "tts": {
                    "charactersCount": summary.tts_characters_count,
                    "audioDuration": summary.tts_audio_duration,
                },
                "stt": { "audioDuration": summary.stt_audio_duration },
            }

        except Exception as e:
            # logger.error(f"Failed to read usage summary: {e}", exc_info=True)
            return

        payload = {
            "conversationId": conversation_id,
            "usage": usage_payload,
            "agentId": agent_id,
        }

        # ---- SEND TO BACKEND ----
        try:
            async with aiohttp.ClientSession() as session:
                async with session.put(
                    f"{backend_url}/v1/conversation/end",
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=10)
                ) as resp:  
                    response_text = await resp.text()

                    if resp.status == 200:
                        # logger.info("Successfully sent end-call usage summary to backend")
                        pass
                    else:
                        # logger.warning(
                        #     f"Backend returned status {resp.status}: {response_text}"
                        # )
                        pass

        except Exception as e:
            # logger.error(f"Error sending end-call usage to backend: {e}", exc_info=True)
            pass
        
    async def shutdown_handler():
        try:
            await send_end_call_to_backend()
        except Exception as e:
            # logger.error(f"Failed to send end-call usage: {e}")
            pass

    ctx.add_shutdown_callback(shutdown_handler)
      
    # Conditionally add avatar based on agent configuration
    if config.get("avatar_enabled", False) and config.get("avatar_id"):  
        try:
            avatar = bey.AvatarSession(avatar_id=config["avatar_id"])  
            await avatar.start(session, room=ctx.room)
            # logger.info(f"Avatar started for agent: {config['agent_name']}")
        except Exception as e:
            # logger.warning(f"Failed to start avatar: {e}")
            pass
    
    # Start the session with the created agent
    await session.start(
        agent=agent,
        room=ctx.room,
        room_options=room_io.RoomOptions(
            audio_input=room_io.AudioInputOptions(
                noise_cancellation=lambda params: noise_cancellation.BVCTelephony()
                if params.participant.kind == rtc.ParticipantKind.PARTICIPANT_KIND_SIP
                else noise_cancellation.BVC(),
            ),
        ),
    )


if __name__ == "__main__":
    cli.run_app(WorkerOptions(entrypoint_fnc=entrypoint, prewarm_fnc=prewarm))