import boto3
import logging
import os
import json
import time
from datetime import datetime
from typing import Dict, Any, Optional, List
from supabase import create_client, Client
import re

from utils.rekognition_utils import check_training_job_status
from utils.logging_utils import setup_logger, reduce_logging_verbosity

# Initialize logger
logger = setup_logger()
reduce_logging_verbosity()

# Initialize clients
rekognition_client = boto3.client('rekognition')
ssm_client = boto3.client('ssm')

# Initialize Supabase client
supabase: Optional[Client] = None

def get_ssm_parameter(param_name):
    """Get a parameter from SSM Parameter Store"""
    try:
        response = ssm_client.get_parameter(Name=param_name, WithDecryption=True)
        return response['Parameter']['Value']
    except Exception as e:
        logger.warning(f"Failed to get SSM parameter {param_name}: {str(e)}")
        return None

def initialize_supabase():
    """Initialize Supabase client using SSM parameters"""
    global supabase
    
    try:
        supabase_url = get_ssm_parameter('/supabase/url')
        supabase_key = get_ssm_parameter('/supabase/anon')
        
        if supabase_url and supabase_key:
            supabase = create_client(supabase_url, supabase_key)
            logger.info("Supabase client initialized using SSM parameters")
            return True
        else:
            logger.warning("Supabase credentials not found in SSM, metadata storage will be skipped")
            return False
    except Exception as e:
        logger.error(f"Failed to initialize Supabase client: {str(e)}")
        return False

def get_training_models_from_supabase() -> List[Dict[str, Any]]:
    """
    Get list of training models from Supabase that are in progress
    
    Returns:
        List of model records with status='TRAINING'
    """
    try:
        if not supabase:
            logger.error("Supabase client not initialized")
            return []
        
        logger.info("Connected to Supabase, querying model_versions table")
        
        # Use a direct table query approach similar to the shared example
        models_query = supabase.table('model_versions').select('*').execute()
        
        if models_query.data is None:
            logger.warning("No data returned from model_versions table query")
            return []
            
        logger.info(f"Total records in model_versions table: {len(models_query.data)}")
        
        if len(models_query.data) > 0:
            # Log a sample record for debugging
            logger.info(f"Sample record: {models_query.data[0]}")
            
            # Log all statuses that exist in the table
            statuses = set(record.get('status') for record in models_query.data if record.get('status'))
            logger.info(f"Statuses found in table: {statuses}")
        
        # Filter for training status records - check for both TRAINING and TRAINING_IN_PROGRESS
        training_models = [model for model in models_query.data if model.get('status') in ['TRAINING', 'TRAINING_IN_PROGRESS']]
        
        if training_models:
            logger.info(f"Found {len(training_models)} models in training")
            for model in training_models:
                logger.info(f"Training model: {model.get('version_name')} - Status: {model.get('status')}")
            return training_models
        else:
            logger.info("No models currently in training")
            return []
            
    except Exception as e:
        logger.error(f"Error querying Supabase for training models: {str(e)}")
        return []

def update_model_status_in_supabase(model_record: Dict[str, Any], status_info: Dict[str, Any]) -> bool:
    """
    Update the status of a model in Supabase
    
    Args:
        model_record: Model record from Supabase
        status_info: Status information from AWS Rekognition
        
    Returns:
        bool: Success flag
    """
    try:
        if not supabase:
            logger.error("Supabase client not initialized")
            return False
        
        # Extract the project_version_arn to use as primary key
        project_version_arn = model_record.get('project_version_arn')
        if not project_version_arn:
            logger.error(f"No project_version_arn found in model record: {model_record}")
            return False
        
        version_name = model_record.get('version_name', 'unknown')
        logger.info(f"Updating model status for {version_name} to {status_info.get('status')}")
        
        # Prepare update data
        update_data = {
            'status': status_info.get('status'),
            'status_message': status_info.get('status_message'),
            'updated_at': datetime.utcnow().isoformat()
        }
        
        # Add evaluation metrics if available
        if 'evaluation' in status_info and status_info['evaluation']:
            metrics_json = json.dumps(status_info['evaluation'])
            update_data['evaluation'] = metrics_json  # Match the field name in your db table
            
            # Check F1 score
            f1_score = status_info['evaluation'].get('F1Score', 0.0)
            logger.info(f"Model F1 score: {f1_score}")
            update_data['f1_score'] = f1_score
            
            # Set meets_criteria based on minimum F1 score (0.85 is a common threshold)
            update_data['meets_criteria'] = f1_score >= 0.85
        
        # Update the record using project_version_arn as the primary key
        response = supabase.table('model_versions').update(update_data).eq('project_version_arn', project_version_arn).execute()
        
        if response and response.data:
            logger.info(f"Successfully updated model status in Supabase for ARN: {project_version_arn}")
            if len(response.data) > 0:
                logger.info(f"Updated record: {response.data[0]}")
            return True
        else:
            # Check if we received any error
            error_msg = getattr(response, 'error', None)
            logger.warning(f"Supabase update may have failed. Error: {error_msg}")
            
            # Try to log the response in a safe way
            if hasattr(response, 'data'):
                logger.info(f"Response data length: {len(response.data) if response.data else 0}")
            else:
                logger.warning("Response has no data attribute")
                
            return False
            
    except Exception as e:
        logger.error(f"Error updating model status in Supabase: {str(e)}")
        return False

def check_and_update_model(model_record: Dict[str, Any]) -> Dict[str, Any]:
    """
    Check the status of a model and update Supabase only if training is complete
    
    Args:
        model_record: Model record from Supabase
        
    Returns:
        Dict with updated status information
    """
    try:
        # Extract needed information
        project_arn = model_record.get('project_arn')
        version_name = model_record.get('version_name')
        
        if not project_arn or not version_name:
            logger.error(f"Missing project_arn or version_name in model record: {model_record}")
            return {
                'success': False,
                'error': 'Missing required fields in model record',
                'version_name': version_name,
                'status': model_record.get('status', 'UNKNOWN')
            }
        
        # Check status in AWS Rekognition
        status_info = check_training_job_status(project_arn, version_name)
        
        # Only update Supabase if training is not in progress
        # Check for both TRAINING_IN_PROGRESS (AWS status) and TRAINING (our database status)
        if status_info.get('status') not in ['TRAINING_IN_PROGRESS', 'TRAINING']:
            # Training is complete (success or failure), update the database
            update_success = update_model_status_in_supabase(model_record, status_info)
            logger.info(f"Training complete for {version_name} with status {status_info.get('status')} - Updated Supabase: {update_success}")
        else:
            # Still in training, don't update the database
            logger.info(f"Model {version_name} is still in training with status {status_info.get('status')} - not updating Supabase")
            update_success = True  # Consider it a success since we're intentionally not updating
        
        return {
            'success': update_success,
            'version_name': version_name,
            'status': status_info.get('status'),
            'status_message': status_info.get('status_message')
        }
        
    except Exception as e:
        logger.error(f"Error checking and updating model: {str(e)}")
        return {
            'success': False,
            'error': str(e),
            'version_name': model_record.get('version_name'),
            'status': model_record.get('status', 'UNKNOWN')
        }

def fix_project_arn(project_arn):
    """
    Fix the project ARN format by ensuring it ends with the correct project ID.
    AWS Rekognition project ARNs must use the correct project ID: 1725357683732
    """
    if not project_arn:
        logger.error("No project ARN provided")
        return None
        
    logger.info(f"Checking project ARN format: {project_arn}")
    
    # Use the known working project ID 
    project_id = "1725357683732"
        
    # Check if the ARN already has a version number
    pattern = r'project\/[a-zA-Z0-9_\.-]{1,255}\/[0-9]+$'
    if re.search(pattern, project_arn):
        # If it matches the pattern but has a different ID, replace it
        if not project_arn.endswith(f"/{project_id}"):
            # Extract everything up to the last slash
            base_arn = '/'.join(project_arn.split('/')[:-1])
            fixed_arn = f"{base_arn}/{project_id}"
            logger.info(f"Updated project ID: {project_arn} -> {fixed_arn}")
            return fixed_arn
        else:
            logger.info(f"Project ARN format is already correct: {project_arn}")
            return project_arn
    
    # Get just the project name without any trailing slashes
    if '/project/' in project_arn:
        base_arn = project_arn.rstrip('/')
        # Standard case: append the correct project ID
        fixed_arn = f"{base_arn}/{project_id}"
        logger.info(f"Fixed ARN by adding correct project ID: {project_arn} -> {fixed_arn}")
        return fixed_arn
    
    logger.warning(f"Cannot determine correct ARN format from: {project_arn}")
    
    # Fall back to simple approach if more sophisticated extraction fails
    if project_arn.endswith('/'):
        fixed_arn = f"{project_arn}{project_id}"
    else:
        fixed_arn = f"{project_arn}/{project_id}"
        
    logger.info(f"Applied fallback ARN fix: {project_arn} -> {fixed_arn}")
    return fixed_arn

def get_project_details(project_arn):
    """
    Get project details including model versions directly from AWS.
    
    Args:
        project_arn: ARN of the project
        
    Returns:
        dict: Project details including models
    """
    try:
        # Extract project name from ARN
        match = re.match(r'arn:aws:rekognition:[^:]+:[^:]+:project/([^/]+)', project_arn)
        if not match:
            logger.error(f"Could not extract project name from ARN: {project_arn}")
            return None
            
        project_name = match.group(1)
        logger.info(f"Looking up project details for: {project_name}")
        
        # Use AWS's recommended approach to get project information
        response = rekognition_client.describe_projects(ProjectNames=[project_name])
        
        if 'ProjectDescriptions' in response and response['ProjectDescriptions']:
            project = response['ProjectDescriptions'][0]
            
            # Now get model versions for this project
            project_versions = []
            try:
                paginator = rekognition_client.get_paginator('describe_project_versions')
                for page in paginator.paginate(ProjectArn=project_arn):
                    project_versions.extend(page.get('ProjectVersionDescriptions', []))
                
                # Add models to project details
                project['ModelVersions'] = project_versions
                logger.info(f"Found {len(project_versions)} model versions for project {project_name}")
            except Exception as e:
                logger.warning(f"Error getting model versions: {str(e)}")
                project['ModelVersions'] = []
                
            return project
        else:
            logger.warning(f"No project found with name: {project_name}")
            return None
    except Exception as e:
        logger.error(f"Error getting project details: {str(e)}")
        return None

def lambda_handler(event, context):
    """
    Lambda handler to check status of training models and update Supabase
    
    Args:
        event: Lambda event
        context: Lambda context
        
    Returns:
        Response with processing results and a training_complete flag for Step Functions
    """
    try:
        logger.info(f"Starting check of training models status")
        
        # Get environment variables
        project_arn = os.environ.get('PROJECT_ARN')
        
        if not project_arn:
            missing = ['PROJECT_ARN']
            
            # Try to get from SSM if not in environment
            if 'PROJECT_ARN' in missing:
                project_name = get_ssm_parameter('/datafy-rekognition-stack/project-name')
                if project_name:
                    # Construct the ARN using the project name with fallback values if needed
                    aws_region = os.environ.get('AWS_REGION', 'eu-west-1')
                    aws_account_id = os.environ.get('AWS_ACCOUNT_ID', '587594388832')
                    project_arn = f"arn:aws:rekognition:{aws_region}:{aws_account_id}:project/{project_name}"
                    logger.info(f"Constructed project ARN from SSM: {project_arn}")
                    missing.remove('PROJECT_ARN')
            
            if missing:
                logger.warning(f"Missing required configuration: {', '.join(missing)}")
        
        # Fix project ARN format
        project_arn = fix_project_arn(project_arn)
        if not project_arn:
            logger.warning("Failed to fix project ARN, will try alternatives")
            project_arn = None
        else:
            logger.info(f"Using fixed project ARN: {project_arn}")
        
        # If we still don't have a working ARN, try using the describe_projects method
        if not project_arn:
            # Try different variations of the project name if provided
            project_names_to_try = []
            
            if project_name:
                # Original project name
                project_names_to_try.append(project_name)
                
                # Try without hyphens
                if '-' in project_name:
                    project_names_to_try.append(project_name.replace('-', ''))
                
                # Try with "Datafy" appended if not already present
                if not project_name.endswith('Datafy'):
                    project_names_to_try.append(f"{project_name}Datafy")
                    project_names_to_try.append("Datafy")  # Just try "Datafy" alone
            else:
                # Fallbacks if no project name provided
                project_names_to_try = ["Datafy", "datafy-rekognition", "datafyrekognition"]
                logger.warning(f"No project name available, will try these names: {project_names_to_try}")
            
            # Try to find the project using the names
            project_found = False
            
            for name in project_names_to_try:
                try:
                    logger.info(f"Trying to find project with name: {name}")
                    project_response = rekognition_client.describe_projects(ProjectNames=[name])
                    
                    if project_response.get('ProjectDescriptions'):
                        raw_project_arn = project_response['ProjectDescriptions'][0]['ProjectArn']
                        logger.info(f"Found project ARN: {raw_project_arn}")
                        
                        # Fix the ARN format
                        project_arn = fix_project_arn(raw_project_arn)
                        if project_arn:
                            logger.info(f"Using fixed project ARN: {project_arn}")
                            project_found = True
                            break
                except Exception as e:
                    logger.warning(f"Could not find project with name '{name}': {str(e)}")
                    continue
            
            # If no project found through names, try listing all projects
            if not project_found:
                try:
                    logger.info("No project found by name, listing all projects")
                    list_response = rekognition_client.describe_projects()
                    
                    if list_response.get('ProjectDescriptions'):
                        # Log all available projects to help with debugging
                        for idx, project in enumerate(list_response['ProjectDescriptions']):
                            logger.info(f"Available project {idx+1}: {project.get('ProjectName')} - {project.get('ProjectArn')}")
                        
                        # Use the first project found
                        if list_response['ProjectDescriptions']:
                            raw_project_arn = list_response['ProjectDescriptions'][0]['ProjectArn']
                            logger.info(f"Using first available project ARN: {raw_project_arn}")
                            
                            # Fix the ARN format
                            project_arn = fix_project_arn(raw_project_arn)
                            if project_arn:
                                logger.info(f"Using fixed project ARN: {project_arn}")
                                project_found = True
                except Exception as e:
                    logger.error(f"Error listing projects: {str(e)}")
        
        if not project_arn:
            logger.error("Could not find any Rekognition projects")
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'message': 'No Rekognition projects found',
                    'error': 'Project not found'
                }),
                'training_complete': False
            }
        
        # Check for any models in TRAINING directly from Rekognition
        try:
            # Get project details including model versions
            project_details = get_project_details(project_arn)
            
            if not project_details:
                logger.warning("Could not get project details, will continue with Supabase check")
            else:
                # Check if any models are in TRAINING status
                model_versions = project_details.get('ModelVersions', [])
                training_models_rekognition = [
                    model for model in model_versions 
                    if model.get('Status') in ['TRAINING_IN_PROGRESS', 'TRAINING']
                ]
                
                if training_models_rekognition:
                    logger.info(f"Found {len(training_models_rekognition)} models in training status directly from Rekognition")
                    for model in training_models_rekognition:
                        logger.info(f"Training model from Rekognition: {model.get('VersionName')} - Status: {model.get('Status')}")
                    
                    # If we find models in training directly from Rekognition, we know training is not complete
                    return {
                        'statusCode': 200,
                        'body': json.dumps({
                            'message': f'Found {len(training_models_rekognition)} models in training from Rekognition',
                            'training_models': [model.get('VersionName') for model in training_models_rekognition]
                        }),
                        'training_complete': False
                    }
                
                logger.info("No models found in training status directly from Rekognition")
        except Exception as e:
            logger.error(f"Error checking model status from Rekognition: {str(e)}")
            # Continue with Supabase check, but log the error
        
        # Initialize Supabase
        if not initialize_supabase():
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': 'Failed to initialize Supabase client',
                    'error': 'Supabase initialization failed'
                }),
                'training_complete': False
            }
        
        # Get models in training from Supabase
        training_models = get_training_models_from_supabase()
        
        if not training_models:
            # If we're checking as part of step function flow, consider as complete
            # since there are no models in TRAINING status in either Rekognition or Supabase
            is_step_function = event.get('is_step_function', False)
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'No models currently in training',
                    'count': 0
                }),
                'training_complete': is_step_function
            }
        
        # Check and update each model
        results = []
        training_complete = True  # Assume training is complete until we find a model still in training
        
        for model in training_models:
            result = check_and_update_model(model)
            results.append(result)
            
            # If any model has a training status, mark as not complete
            if result.get('status') in ['TRAINING_IN_PROGRESS', 'TRAINING']:
                training_complete = False
        
        # Return results with training_complete flag for Step Functions
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': f'Processed {len(results)} models',
                'results': results
            }),
            'training_complete': training_complete
        }
        
    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}", exc_info=True)
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': f'Error processing models: {str(e)}',
                'error': 'Internal server error'
            }),
            'training_complete': False
        }