# backend/lib/rekognition_training.py

import boto3
import json
import time
from botocore.exceptions import ClientError

class RekognitionTraining:
    def __init__(self, project_name, version_name):
        self.rekognition = boto3.client('rekognition')
        self.project_name = project_name
        self.version_name = version_name

    async def start_model_training(self):
        """Inicia el entrenamiento del modelo"""
        try:
            response = self.rekognition.create_project_version(
                ProjectArn=self.project_name,
                VersionName=self.version_name,
                OutputConfig={
                    'S3Bucket': self.training_bucket,
                    'S3KeyPrefix': f'output/{self.version_name}'
                },
                TrainingParameters={
                    'MAX_ITERATIONS': 100,
                    'LEARNING_RATE': '0.001',
                    'AUTO_STOP': True
                }
            )
            
            return response['ProjectVersionArn']
        except ClientError as e:
            print(f"Error starting model training: {e}")
            raise

    async def check_training_status(self, version_arn):
        """Verifica el estado del entrenamiento"""
        try:
            while True:
                response = self.rekognition.describe_project_versions(
                    ProjectArn=self.project_name,
                    VersionNames=[self.version_name]
                )
                
                status = response['ProjectVersionDescriptions'][0]['Status']
                
                if status == 'TRAINING_COMPLETED':
                    return True
                elif status in ['TRAINING_FAILED', 'FAILED']:
                    raise Exception(f"Training failed: {response['ProjectVersionDescriptions'][0]['StatusMessage']}")
                
                await asyncio.sleep(60)  # Check every minute
                
        except ClientError as e:
            print(f"Error checking training status: {e}")
            raise

    async def start_model(self, version_arn):
        """Inicia el modelo entrenado"""
        try:
            response = self.rekognition.start_project_version(
                ProjectVersionArn=version_arn,
                MinInferenceUnits=1
            )
            return response
        except ClientError as e:
            print(f"Error starting model: {e}")
            raise

    async def evaluate_model(self, version_arn):
        """Evalúa el rendimiento del modelo"""
        try:
            response = self.rekognition.evaluate_project_version(
                ProjectVersionArn=version_arn
            )
            
            evaluation_results = {
                'f1_score': response['EvaluationResults']['F1Score'],
                'precision': response['EvaluationResults']['Precision'],
                'recall': response['EvaluationResults']['Recall'],
                'confidence_threshold': response['EvaluationResults']['ConfidenceThreshold']
            }
            
            return evaluation_results
        except ClientError as e:
            print(f"Error evaluating model: {e}")
            raise

    def stop_model(self, version_arn):
        """Detiene el modelo"""
        try:
            response = self.rekognition.stop_project_version(
                ProjectVersionArn=version_arn
            )
            return response
        except ClientError as e:
            print(f"Error stopping model: {e}")
            raise

# backend/lib/image_processing.py

import boto3
import json
from datetime import datetime
from typing import Dict, List, Optional

class ImageProcessor:
    def __init__(self):
        self.rekognition = boto3.client('rekognition')
        self.s3 = boto3.client('s3')
        self.dynamodb = boto3.resource('dynamodb')

    async def process_medical_image(
        self,
        bucket: str,
        key: str,
        metadata: Dict,
        model_arn: str
    ) -> Dict:
        """Procesa una imagen médica usando Rekognition Custom Labels"""
        try:
            # Detectar etiquetas personalizadas
            custom_labels = await self._detect_custom_labels(bucket, key, model_arn)
            
            # Análisis de calidad de imagen
            quality_analysis = await self._analyze_image_quality(bucket, key)
            
            # Metadata enriquecida
            enriched_metadata = self._enrich_metadata(metadata, custom_labels, quality_analysis)
            
            # Guardar resultados
            analysis_id = await self._save_analysis_results(
                bucket,
                key,
                enriched_metadata,
                custom_labels,
                quality_analysis
            )
            
            return {
                'analysis_id': analysis_id,
                'custom_labels': custom_labels,
                'quality_analysis': quality_analysis,
                'metadata': enriched_metadata
            }
            
        except Exception as e:
            print(f"Error processing image: {e}")
            raise

    async def _detect_custom_labels(
        self,
        bucket: str,
        key: str,
        model_arn: str
    ) -> List[Dict]:
        """Detecta etiquetas personalizadas en la imagen"""
        try:
            response = self.rekognition.detect_custom_labels(
                Image={
                    'S3Object': {
                        'Bucket': bucket,
                        'Name': key
                    }
                },
                ProjectVersionArn=model_arn,
                MinConfidence=90.0
            )
            
            return [{
                'name': label['Name'],
                'confidence': label['Confidence'],
                'coordinates': label.get('Geometry', {})
            } for label in response['CustomLabels']]
            
        except Exception as e:
            print(f"Error detecting custom labels: {e}")
            raise

    async def _analyze_image_quality(
        self,
        bucket: str,
        key: str
    ) -> Dict:
        """Analiza la calidad de la imagen"""
        try:
            # Obtener metadata de S3
            image_metadata = self.s3.head_object(Bucket=bucket, Key=key)
            
            # Analizar calidad
            quality_analysis = {
                'format': key.split('.')[-1].upper(),
                'size_bytes': image_metadata['ContentLength'],
                'content_type': image_metadata['ContentType'],
                'timestamp': image_metadata['LastModified'].isoformat()
            }
            
            # Análisis adicional con Rekognition
            quality_response = self.rekognition.detect_moderation_labels(
                Image={
                    'S3Object': {
                        'Bucket': bucket,
                        'Name': key
                    }
                }
            )
            
            quality_analysis['moderation'] = {
                'issues_detected': len(quality_response['ModerationLabels']) == 0,
                'confidence': quality_response['ModerationLabels'][0]['Confidence'] if quality_response['ModerationLabels'] else 100
            }
            
            return quality_analysis
            
        except Exception as e:
            print(f"Error analyzing image quality: {e}")
            raise

    def _enrich_metadata(
        self,
        original_metadata: Dict,
        custom_labels: List[Dict],
        quality_analysis: Dict
    ) -> Dict:
        """Enriquece los metadatos con los resultados del análisis"""
        return {
            **original_metadata,
            'analysis_timestamp': datetime.utcnow().isoformat(),
            'detected_conditions': [label['name'] for label in custom_labels],
            'quality_metrics': {
                'format': quality_analysis['format'],
                'size': quality_analysis['size_bytes'],
                'quality_score': quality_analysis['moderation']['confidence']
            }
        }

    async def _save_analysis_results(
        self,
        bucket: str,
        key: str,
        metadata: Dict,
        custom_labels: List[Dict],
        quality_analysis: Dict
    ) -> str:
        """Guarda los resultados del análisis en DynamoDB"""
        try:
            table = self.dynamodb.Table(os.environ['ANALYSES_TABLE'])
            
            analysis_id = f"analysis_{int(datetime.utcnow().timestamp())}"
            
            item = {
                'analysis_id': analysis_id,
                'image_bucket': bucket,
                'image_key': key,
                'timestamp': int(datetime.utcnow().timestamp()),
                'metadata': metadata,
                'custom_labels': custom_labels,
                'quality_analysis': quality_analysis,
                'status': 'COMPLETED'
            }
            
            table.put_item(Item=item)
            
            return analysis_id
            
        except Exception as e:
            print(f"Error saving analysis results: {e}")
            raise

# backend/functions/rekognition/handlers.py

import json
import os
from typing import Dict, Any
from lib.image_processing import ImageProcessor
from lib.rekognition_training import RekognitionTraining
from lib.models import AnalysisRequest, TrainingRequest, ModelStatus

class RekognitionAPI:
    def __init__(self):
        self.processor = ImageProcessor()
        self.trainer = RekognitionTraining(
            project_name=os.environ['REKOGNITION_PROJECT_NAME'],
            version_name=f"v{int(datetime.utcnow().timestamp())}"
        )

    async def handle_image_analysis(self, event: Dict[str, Any]) -> Dict:
        """Maneja el análisis de imágenes médicas"""
        try:
            body = json.loads(event['body'])
            request = AnalysisRequest(**body)

            # Validar imagen
            if not await self._validate_image(request.bucket, request.key):
                return self._error_response(400, 'Invalid image format or size')

            # Procesar imagen
            results = await self.processor.process_medical_image(
                bucket=request.bucket,
                key=request.key,
                metadata=request.metadata,
                model_arn=os.environ['REKOGNITION_MODEL_ARN']
            )

            return self._success_response(results)

        except Exception as e:
            return self._error_response(500, str(e))

    async def handle_model_training(self, event: Dict[str, Any]) -> Dict:
        """Maneja el entrenamiento del modelo"""
        try:
            body = json.loads(event['body'])
            request = TrainingRequest(**body)

            # Iniciar entrenamiento
            version_arn = await self.trainer.start_model_training()

            # Verificar estado
            if not await self.trainer.check_training_status(version_arn):
                return self._error_response(500, 'Model training failed')

            # Evaluar modelo
            evaluation_results = await self.trainer.evaluate_model(version_arn)

            # Iniciar modelo si la evaluación es satisfactoria
            if evaluation_results['f1_score'] >= 0.90:
                await self.trainer.start_model(version_arn)
                
                return self._success_response({
                    'version_arn': version_arn,
                    'evaluation': evaluation_results,
                    'status': 'DEPLOYED'
                })
            else:
                return self._success_response({
                    'version_arn': version_arn,
                    'evaluation': evaluation_results,
                    'status': 'EVALUATION_FAILED'
                })

        except Exception as e:
            return self._error_response(500, str(e))

    async def handle_model_status(self, event: Dict[str, Any]) -> Dict:
        """Obtiene el estado del modelo"""
        try:
            model_arn = event['pathParameters']['modelArn']
            status = await self.trainer.get_model_status(model_arn)
            
            return self._success_response({
                'model_arn': model_arn,
                'status': status.value,
                'details': await self.trainer.get_model_details(model_arn)
            })

        except Exception as e:
            return self._error_response(500, str(e))

    async def handle_batch_analysis(self, event: Dict[str, Any]) -> Dict:
        """Maneja el análisis por lotes de imágenes"""
        try:
            body = json.loads(event['body'])
            images = body['images']
            
            results = []
            for image in images:
                result = await self.processor.process_medical_image(
                    bucket=image['bucket'],
                    key=image['key'],
                    metadata=image.get('metadata', {}),
                    model_arn=os.environ['REKOGNITION_MODEL_ARN']
                )
                results.append(result)

            return self._success_response({
                'batch_id': f"batch_{int(datetime.utcnow().timestamp())}",
                'results': results,
                'summary': self._generate_batch_summary(results)
            })

        except Exception as e:
            return self._error_response(500, str(e))

    async def handle_model_metrics(self, event: Dict[str, Any]) -> Dict:
        """Obtiene métricas del modelo"""
        try:
            model_arn = event['pathParameters']['modelArn']
            time_range = event.get('queryStringParameters', {}).get('timeRange', '24h')
            
            metrics = await self.trainer.get_model_metrics(model_arn, time_range)
            
            return self._success_response({
                'model_arn': model_arn,
                'time_range': time_range,
                'metrics': metrics
            })

        except Exception as e:
            return self._error_response(500, str(e))

    async def handle_custom_labels_management(self, event: Dict[str, Any]) -> Dict:
        """Gestiona las etiquetas personalizadas"""
        try:
            operation = event['httpMethod']
            
            if operation == 'GET':
                labels = await self.trainer.get_custom_labels()
                return self._success_response({'labels': labels})
                
            elif operation == 'POST':
                body = json.loads(event['body'])
                new_label = await self.trainer.add_custom_label(
                    name=body['name'],
                    description=body.get('description', ''),
                    metadata=body.get('metadata', {})
                )
                return self._success_response({'label': new_label})
                
            elif operation == 'DELETE':
                label_name = event['pathParameters']['labelName']
                await self.trainer.delete_custom_label(label_name)
                return self._success_response({
                    'message': f'Label {label_name} deleted successfully'
                })

        except Exception as e:
            return self._error_response(500, str(e))

    def _generate_batch_summary(self, results: List[Dict]) -> Dict:
        """Genera un resumen del análisis por lotes"""
        total_images = len(results)
        successful = len([r for r in results if r.get('status') == 'SUCCESS'])
        failed = total_images - successful
        
        conditions_found = {}
        for result in results:
            for label in result.get('custom_labels', []):
                conditions_found[label['name']] = conditions_found.get(label['name'], 0) + 1

        return {
            'total_images': total_images,
            'successful_analyses': successful,
            'failed_analyses': failed,
            'conditions_found': conditions_found,
            'average_confidence': sum(r.get('confidence', 0) for r in results) / total_images
        }

    def _success_response(self, body: Dict) -> Dict:
        """Genera una respuesta exitosa"""
        return {
            'statusCode': 200,
            'headers': {
                'Content-Type': 'application/json',
                'Access-Control-Allow-Origin': '*'
            },
            'body': json.dumps(body)
        }

    def _error_response(self, status_code: int, message: str) -> Dict:
        """Genera una respuesta de error"""
        return {
            'statusCode': status_code,
            'headers': {
                'Content-Type': 'application/json',
                'Access-Control-Allow-Origin': '*'
            },
            'body': json.dumps({
                'error': message
            })
        }

# Rutas disponibles:
# POST /analyze - Analiza una imagen individual
# POST /analyze/batch - Analiza un lote de imágenes
# POST /models/train - Inicia el entrenamiento de un nuevo modelo
# GET /models/{modelArn}/status - Obtiene el estado de un modelo
# GET /models/{modelArn}/metrics - Obtiene métricas de un modelo
# GET /labels - Obtiene etiquetas personalizadas
# POST /labels - Crea una nueva etiqueta personalizada
# DELETE /labels/{labelName} - Elimina una etiqueta personalizada