import os
import logging
import json
import time
from datetime import datetime
from typing import Dict, Any, Optional

from utils.rekognition_utils import check_training_job_status

logger = logging.getLogger(__name__)

def store_training_metadata(
    training_info: Dict[str, Any],
    supabase_client: Optional[Any] = None
) -> None:
    """
    Store training job metadata in Supabase.
    
    Args:
        training_info: Dictionary containing training information
        supabase_client: Optional Supabase client
    """
    try:
        if not supabase_client:
            logger.info("No Supabase client provided, skipping metadata storage")
            return
            
        # Prepare metadata
        metadata = {
            'version_name': training_info.get('version_name'),
            'training_timestamp': int(time.time()),
            'project_version_arn': training_info.get('project_version_arn'),
            'project_arn': training_info.get('project_arn'),
            'status': training_info.get('status'),
            'status_message': training_info.get('status_message'),
            'manifest_file': training_info.get('manifest_file'),
            'bucket': training_info.get('bucket'),
            'created_at': datetime.utcnow().isoformat()
        }
        
        # Log the metadata we're storing to help with troubleshooting
        logger.info(f"Storing metadata to Supabase: {json.dumps(metadata, default=str)}")
        
        # Store in Supabase
        response = supabase_client.table('model_versions').insert(metadata).execute()
        logger.info(f"Stored metadata for version: {metadata['version_name']}")
        
    except Exception as e:
        logger.error(f"Error storing training metadata: {str(e)}")
        # Don't raise the error as this is not critical for training

def poll_training_status(
    project_arn: str,
    version_name: str,
    max_attempts: int = 12,
    delay_seconds: int = 300
) -> Dict[str, Any]:
    """
    Poll the status of a training job until it completes or fails.
    
    Args:
        project_arn: The project ARN
        version_name: The version name to check
        max_attempts: Maximum number of polling attempts
        delay_seconds: Delay between polling attempts in seconds
        
    Returns:
        Dictionary containing final status information
    """
    try:
        logger.info(f"Polling training status for version: {version_name}")
        
        for attempt in range(max_attempts):
            # Check status
            status_info = check_training_job_status(project_arn, version_name)
            status = status_info['status']
            
            # Log progress
            logger.info(f"Training status ({attempt + 1}/{max_attempts}): {status}")
            if 'status_message' in status_info:
                logger.info(f"Status message: {status_info['status_message']}")
                
            # Check if training is complete
            if status in ['RUNNING', 'TRAINING']:
                # Wait before next attempt
                time.sleep(delay_seconds)
            else:
                # Training completed or failed
                return status_info
                
        # Max attempts reached
        logger.warning(f"Max polling attempts ({max_attempts}) reached")
        return {
            'status': 'TIMEOUT',
            'version_name': version_name,
            'message': f"Polling timed out after {max_attempts} attempts"
        }
        
    except Exception as e:
        logger.error(f"Error polling training status: {str(e)}")
        raise

def evaluate_model(
    project_arn: str,
    version_name: str,
    min_f1_score: float = 0.85
) -> Dict[str, Any]:
    """
    Evaluate the model using F1 score and other metrics.
    
    Args:
        project_arn: The project ARN
        version_name: The version name to evaluate
        min_f1_score: Minimum required F1 score (default 0.85)
        
    Returns:
        Dictionary containing evaluation results
    """
    try:
        status_info = check_training_job_status(project_arn, version_name)
        
        evaluation_results = {
            'version_name': version_name,
            'status': status_info.get('status'),
            'meets_criteria': False,
            'metrics': {}
        }
        
        if 'evaluation' in status_info:
            metrics = status_info['evaluation']
            evaluation_results['metrics'] = metrics
            
            # Check F1 score
            f1_score = metrics.get('F1Score', 0.0)
            logger.info(f"Model F1 score: {f1_score}")
            evaluation_results['meets_criteria'] = f1_score >= min_f1_score
            
            # Add evaluation timestamp
            evaluation_results['evaluated_at'] = datetime.utcnow().isoformat()
            
        return evaluation_results
        
    except Exception as e:
        logger.error(f"Error evaluating model: {str(e)}")
        raise