import boto3
import uuid
import json
from datetime import datetime
from botocore.exceptions import ClientError

# Inicializar el cliente de Bedrock
client = boto3.client('bedrock-agent-runtime', region_name='us-west-2')

def create_troubleshooting_session(incident_id, system_affected):
    """
    Crea una nueva sesión para un incidente de infraestructura.
    
    Args:
        incident_id (str): ID del incidente en el sistema de tickets
        system_affected (str): Sistema afectado (ej: "payment-microservice")
        
    Returns:
        str: ID de la sesión creada
    """
    try:
        # Crear una sesión con metadatos relevantes para diagnóstico
        response = client.create_session(
            sessionMetadata={
                "incidentId": incident_id,
                "systemAffected": system_affected,
                "severity": "high",
                "startedAt": datetime.now().isoformat()
            },
            tags={
                'Environment': 'Production',
                'IncidentType': 'PerformanceDegradation'
            }
        )
        
        session_id = response["sessionId"]
        print(f"Sesión de diagnóstico creada. ID: {session_id}")
        return session_id
    
    except ClientError as e:
        print(f"Error al crear la sesión: {str(e)}")
        raise

def store_diagnostic_step(session_identifier, engineer_id, diagnostics_data, screenshots=None):
    """
    Almacena un paso en el proceso de diagnóstico.
    
    Args:
        session_identifier (str): ID o ARN de la sesión
        engineer_id (str): ID del ingeniero ejecutando este paso
        diagnostics_data (dict): Datos del diagnóstico
        screenshots (list, optional): Capturas de pantalla en bytes
    """
    try:
        # Crear una invocación para este paso de diagnóstico
        invocation_id = client.create_invocation(
            sessionIdentifier=session_identifier,
            description=f"Diagnóstico en {diagnostics_data.get('component', 'sistema desconocido')} por {engineer_id}"
        )["invocationId"]
        
        # Estructurar los datos de diagnóstico
        formatted_data = (
            f"## Paso de diagnóstico\n\n"
            f"**Ingeniero:** {engineer_id}\n"
            f"**Componente:** {diagnostics_data.get('component', 'No especificado')}\n"
            f"**Acción ejecutada:** {diagnostics_data.get('action', 'No especificada')}\n\n"
            f"**Resultado observado:**\n{diagnostics_data.get('result', 'No documentado')}\n\n"
            f"**Siguiente acción recomendada:**\n{diagnostics_data.get('next_steps', 'No definida')}"
        )
        
        # Preparar los bloques de contenido
        content_blocks = [
            {
                'text': formatted_data
            }
        ]
        
        # Agregar capturas de pantalla si existen
        if screenshots:
            for i, screenshot in enumerate(screenshots):
                content_blocks.append({
                    'image': {
                        'format': 'png',
                        'source': {'bytes': screenshot}
                    }
                })
        
        # Almacenar el paso de diagnóstico con el parámetro requerido
        client.put_invocation_step(
            sessionIdentifier=session_identifier,
            invocationIdentifier=invocation_id,
            invocationStepId=str(uuid.uuid4()),
            invocationStepTime=datetime.now().isoformat(),  # Este parámetro es obligatorio
            payload={
                'contentBlocks': content_blocks
            }
        )
        
        print(f"✅ Paso de diagnóstico registrado con éxito (invocación: {invocation_id})")
        return invocation_id
    
    except ClientError as e:
        error_code = e.response['Error']['Code'] if 'Error' in e.response and 'Code' in e.response['Error'] else "Desconocido"
        if error_code == 'ThrottlingException':
            print(f"⚠️ Límite de velocidad excedido. Intente nuevamente más tarde.")
        elif error_code == 'ValidationException':
            print(f"❌ Error de validación: {e.response['Error'].get('Message', 'Sin detalle')}")
        else:
            print(f"❌ Error al almacenar el diagnóstico: {str(e)}")
        raise

def retrieve_diagnostic_context(session_identifier):
    """
    Recupera el contexto completo de un diagnóstico de infraestructura.
    
    Args:
        session_identifier (str): ID o ARN de la sesión
        
    Returns:
        dict: Contexto completo del diagnóstico con datos estructurados
    """
    try:
        print("[*] Recuperando contexto de diagnóstico...")
        
        # Obtener detalles de la sesión
        session_response = client.get_session(
            sessionIdentifier=session_identifier
        )
        
        # Manejar diferentes posibles estructuras de respuesta
        if "session" in session_response:
            session = session_response["session"]
        else:
            session = session_response
        
        # Verificar que tenemos acceso a los metadatos
        session_metadata_key = "sessionMetadata"
        if session_metadata_key not in session:
            session_metadata_key = "metadata"  # Alternativa posible
            if session_metadata_key not in session:
                incident_metadata = {}
                print("⚠️ No se pudieron recuperar metadatos de la sesión")
            else:
                incident_metadata = session[session_metadata_key]
        else:
            incident_metadata = session[session_metadata_key]
        
        # Listar todas las invocaciones (pasos de diagnóstico)
        invocations_response = client.list_invocations(
            sessionIdentifier=session_identifier
        )
        
        # CAMBIO CLAVE: Usar invocationSummaries en lugar de invocations
        invocations = invocations_response.get("invocationSummaries", [])
        print(f"[*] Invocaciones recuperadas: {len(invocations)}")
        
        # Construir el contexto estructurado del diagnóstico
        diagnostic_context = {
            "incidentInfo": {
                "incidentId": incident_metadata.get("incidentId", "Unknown"),
                "systemAffected": incident_metadata.get("systemAffected", "Unknown"),
                "severity": incident_metadata.get("severity", "Unknown"),
                "startedAt": session.get("creationDateTime", datetime.now().isoformat()),
                "status": "Active" if not session.get("endDateTime") else "Closed"
            },
            "diagnosticTimeline": [],
            "hypotheses": [],
            "componentsTested": set(),
            "screenshots": []
        }
        
        # Recuperar y organizar los pasos de diagnóstico
        for inv in sorted(invocations, key=lambda x: x.get("createdAt", "")):
            try:
                # Extraer información disponible directamente de la invocación
                invocation_id = inv["invocationId"]
                
                # Convierte createdAt a string ISO si es un objeto datetime
                creation_time = inv.get("createdAt")
                if isinstance(creation_time, datetime):
                    creation_time = creation_time.isoformat()
                
                # La descripción puede no estar disponible
                description = inv.get("description", f"Invocación {invocation_id}")
                
                # Listar pasos de la invocación
                invocation_steps_response = client.list_invocation_steps(
                    sessionIdentifier=session_identifier,
                    invocationIdentifier=invocation_id
                )
                
                # CAMBIO CLAVE: Usar invocationStepSummaries en lugar de invocationSteps
                invocation_steps = invocation_steps_response.get("invocationStepSummaries", [])
                print(f"[*] Pasos encontrados para invocación {invocation_id}: {len(invocation_steps)}")
                
                diagnostic_steps = []
                
                for step in sorted(invocation_steps, key=lambda x: x.get("invocationStepTime", "")):
                    try:
                        step_id = step.get("invocationStepId")
                        
                        # Obtener detalles del paso
                        step_response = client.get_invocation_step(
                            sessionIdentifier=session_identifier,
                            invocationIdentifier=invocation_id,
                            invocationStepId=step_id
                        )
                        
                        # Manejar posibles estructuras diferentes
                        if "invocationStep" in step_response:
                            step_details = step_response["invocationStep"]
                        else:
                            step_details = step_response
                        
                        # Asegurarse de que payload y contentBlocks existen
                        if "payload" not in step_details or "contentBlocks" not in step_details["payload"]:
                            print(f"⚠️ Estructura de paso inesperada para {step_id}")
                            continue
                        
                        # Procesar los bloques de contenido
                        content_blocks = step_details["payload"]["contentBlocks"]
                        text_content = ""
                        images = []
                        
                        for block in content_blocks:
                            if 'text' in block:
                                text_content = block['text']
                                
                                # Extraer componentes probados del texto (lógica más flexible)
                                if "componente:" in text_content.lower() or "Componente:" in text_content:
                                    component = ""
                                    if "Componente:" in text_content:
                                        parts = text_content.split("Componente:")[1].split("\n")
                                        component = parts[0].strip()
                                    elif "componente:" in text_content.lower():
                                        parts = text_content.lower().split("componente:")[1].split("\n")
                                        component = parts[0].strip()
                                    
                                    if component:
                                        diagnostic_context["componentsTested"].add(component)
                                
                                # Extraer hipótesis del texto
                                if "hipótesis" in text_content.lower():
                                    hypothesis_text = text_content
                                    engineer = "Unknown"
                                    if "Ingeniero:" in text_content:
                                        engineer = text_content.split("Ingeniero:")[1].split("\n")[0].strip()
                                    
                                    diagnostic_context["hypotheses"].append({
                                        "text": hypothesis_text,
                                        "timestamp": step_details.get("invocationStepTime", ""),
                                        "engineer": engineer
                                    })
                            
                            if 'image' in block:
                                # Referencia a la imagen
                                images.append({
                                    "stepId": step_id,
                                    "format": block['image'].get('format', 'unknown')
                                })
                                diagnostic_context["screenshots"].append({
                                    "stepId": step_id,
                                    "invocationId": invocation_id,
                                    "timestamp": step_details.get("invocationStepTime", ""),
                                    "associatedText": text_content[:100] + "..." if len(text_content) > 100 else text_content
                                })
                        
                        # Crear entrada para este paso
                        diagnostic_steps.append({
                            "timestamp": step_details.get("invocationStepTime", ""),
                            "textContent": text_content,
                            "hasImages": len(images) > 0,
                            "imageRefs": images
                        })
                    except Exception as step_error:
                        print(f"⚠️ Error procesando paso {step.get('invocationStepId', 'unknown')}: {str(step_error)}")
                        continue
                
                # Extraer ingeniero del descriptor de la invocación (si existe)
                engineer = "Unknown"
                if description and isinstance(description, str) and "por " in description:
                    engineer = description.split("por ")[1]
                
                # Añadir esta invocación al timeline
                diagnostic_context["diagnosticTimeline"].append({
                    "timestamp": creation_time,
                    "description": description,
                    "engineer": engineer,
                    "steps": diagnostic_steps
                })
            except Exception as inv_error:
                print(f"⚠️ Error procesando invocación {inv.get('invocationId', 'unknown')}: {str(inv_error)}")
                continue
        
        # Convertir el conjunto de componentes a lista para serialización JSON
        diagnostic_context["componentsTested"] = list(diagnostic_context["componentsTested"])
        
        # Ordenar cronológicamente
        diagnostic_context["diagnosticTimeline"].sort(key=lambda x: x["timestamp"])
        diagnostic_context["hypotheses"].sort(key=lambda x: x["timestamp"])
        diagnostic_context["screenshots"].sort(key=lambda x: x["timestamp"])
        
        print("✅ Contexto de diagnóstico recuperado correctamente")
        return diagnostic_context
    
    except ClientError as e:
        if e.response['Error']['Code'] == 'ResourceNotFoundException':
            print(f"❌ Error: La sesión {session_identifier} no existe")
        else:
            print(f"❌ Error al recuperar el contexto del diagnóstico: {str(e)}")
        return None
    except Exception as e:
        print(f"❌ Error inesperado: {str(e)}")
        import traceback
        traceback.print_exc()  # Para obtener el stack trace completo
        return None

def end_diagnostic_session(session_identifier, resolution_summary, resolution_type):
    """
    Finaliza una sesión de diagnóstico de infraestructura con información 
    de resolución.
    
    Args:
        session_identifier (str): ID o ARN de la sesión
        resolution_summary (str): Resumen de cómo se resolvió el incidente
        resolution_type (str): Categoría de resolución (fix, workaround, escalation)
    """
    try:
        # Primero, añadimos un paso final con el resumen de resolución
        invocation_id = client.create_invocation(
            sessionIdentifier=session_identifier,
            description="Resolución del incidente"
        )["invocationId"]
        
        # Estructurar el resumen de resolución
        resolution_data = (
            f"## Resolución del Incidente\n\n"
            f"**Tipo de resolución:** {resolution_type}\n\n"
            f"**Resumen:**\n{resolution_summary}\n\n"
            f"**Fecha de resolución:** {datetime.now().isoformat()}\n\n"
            f"**Lecciones aprendidas:**\n- [Por completar en la revisión post-incidente]"
        )
        
        # Almacenar el resumen de resolución - AÑADIDO EL PARÁMETRO FALTANTE
        client.put_invocation_step(
            sessionIdentifier=session_identifier,
            invocationIdentifier=invocation_id,
            invocationStepId=str(uuid.uuid4()),
            invocationStepTime=datetime.now().isoformat(),  # ¡Este parámetro es crucial!
            payload={
                'contentBlocks': [{
                    'text': resolution_data
                }]
            }
        )
        
        # Ahora finalizamos formalmente la sesión
        client.end_session(
            sessionIdentifier=session_identifier
        )
        
        print(f"✅ Sesión de diagnóstico {session_identifier} finalizada con éxito")
        
        # Opcionalmente, podríamos exportar un resumen completo a un sistema de gestión de conocimiento
        export_diagnostic_knowledge(session_identifier, resolution_type)
        
    except ClientError as e:
        print(f"❌ Error al finalizar la sesión de diagnóstico: {str(e)}")
        raise

def delete_diagnostic_session(session_identifier, reason, approver_id):
    """
    Elimina permanentemente una sesión de diagnóstico y todos sus datos asociados.
    
    Args:
        session_identifier (str): ID o ARN de la sesión
        reason (str): Justificación para la eliminación
        approver_id (str): ID del responsable que aprueba la eliminación
    """
    try:
        # Primero, registramos la solicitud de eliminación en sistemas de auditoría
        # (esto podría ser un sistema externo en un escenario real)
        audit_log = {
            "action": "session_deletion",
            "session_id": session_identifier,
            "timestamp": datetime.now().isoformat(),
            "reason": reason,
            "approver": approver_id
        }
        
        print(f"Registrando eliminación en logs de auditoría: {json.dumps(audit_log)}")
        
        # Ahora eliminamos la sesión y todos sus datos asociados
        client.delete_session(
            sessionIdentifier=session_identifier
        )
        
        print(f"Sesión de diagnóstico {session_identifier} eliminada permanentemente")
        
        # En un entorno de producción, podrías notificar a sistemas de monitoreo
        # o a equipos relevantes sobre la eliminación
        
    except ClientError as e:
        print(f"Error al eliminar la sesión de diagnóstico: {str(e)}")
        raise

def create_secure_session():
    try:
        session_id = client.create_session(
            # Usar una clave KMS personalizada
            encryptionKeyArn="arn:aws:kms:us-west-2:123456789012:key/your-key-id"
        )["sessionId"]
        print(f"Sesión segura creada. ID: {session_id}")
        return session_id
    except ClientError as e:
        print(f"Error: {e}")

def store_screenshot(session_id, invocation_id, screenshot_data, metadata):
    """
    Almacena una captura de pantalla con metadatos para facilitar búsquedas.
    """
    step_id = str(uuid.uuid4())
    
    # Almacenar solo los metadatos de la imagen
    client.put_invocation_step(
        sessionIdentifier=session_id,
        invocationIdentifier=invocation_id,
        invocationStepId=f"{step_id}-metadata",
        payload={
            'contentBlocks': [{
                'text': json.dumps(metadata)
            }]
        }
    )
    
    # Almacenar la imagen con referencia a sus metadatos
    client.put_invocation_step(
        sessionIdentifier=session_id,
        invocationIdentifier=invocation_id,
        invocationStepId=step_id,
        payload={
            'contentBlocks': [{
                'image': {
                    'format': 'png',
                    'source': {'bytes': screenshot_data}
                }
            }]
        }
    )