import boto3
import logging
import os
import uuid
import json
import time
import re
from datetime import datetime
from botocore.exceptions import ClientError
import copy

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')
supabase_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 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_supabase_client():
    """Initialize the Supabase client"""
    global supabase_client
    
    if supabase_client is not None:
        return supabase_client
        
    try:
        from supabase import create_client

        # Get Supabase URL and key from environment variables or SSM
        supabase_url = os.environ.get('SUPABASE_URL')
        supabase_key = os.environ.get('SUPABASE_KEY')
        
        # If not in env vars, try SSM
        if not supabase_url:
            supabase_url = get_ssm_parameter('/datafy-rekognition-stack/supabase-url')
        
        if not supabase_key:
            supabase_key = get_ssm_parameter('/datafy-rekognition-stack/supabase-anon-key')
            
        if not supabase_url or not supabase_key:
            logger.error("Supabase credentials not found")
            return None
            
        # Initialize Supabase client
        supabase_client = create_client(supabase_url, supabase_key)
        logger.info("Supabase client initialized")
        return supabase_client
        
    except Exception as e:
        logger.error(f"Failed to initialize Supabase client: {str(e)}")
        return None

def update_supabase_dataset(dataset_id, status, model_id=None, error_message=None):
    """Update a dataset record in Supabase"""
    try:
        supabase = get_supabase_client()
        if not supabase:
            logger.error("No Supabase client available")
            return False
            
        # Prepare update data
        update_data = {
            'status': status,
            'updated_at': datetime.utcnow().isoformat()
        }
        
        # Add model_id if provided
        if model_id:
            update_data['model_id'] = model_id
            
        # Add error message if provided
        if error_message:
            update_data['error_message'] = error_message
            
        # Update the dataset in Supabase
        result = supabase.table('datasets').update(update_data).eq('id', dataset_id).execute()
        
        if result.data:
            logger.info(f"Updated dataset {dataset_id} status to {status}")
            return True
        else:
            logger.warning(f"No dataset updated, check if dataset {dataset_id} exists")
            return False
            
    except Exception as e:
        logger.error(f"Error updating dataset in Supabase: {str(e)}")
        return False

def save_model_to_supabase(project_version_arn, dataset_id, model_name, user_id):
    """Save model information to Supabase"""
    try:
        supabase = get_supabase_client()
        if not supabase:
            logger.error("No Supabase client available")
            return None
            
        # Generate a new UUID for the model
        model_id = str(uuid.uuid4())
        
        # Create model record
        model_data = {
            'id': model_id,
            'dataset_id': dataset_id,
            'name': model_name,
            'arn': project_version_arn,
            'status': 'TRAINING',
            'created_by': user_id,
            'created_at': datetime.utcnow().isoformat(),
            'updated_at': datetime.utcnow().isoformat()
        }
        
        # Insert the model into Supabase
        result = supabase.table('models').insert(model_data).execute()
        
        if result.data:
            logger.info(f"Created model record with ID: {model_id}")
            return model_id
        else:
            logger.warning("No model record created")
            return None
            
    except Exception as e:
        logger.error(f"Error saving model to Supabase: {str(e)}")
        return None

def create_model(project_arn, dataset_arn, model_name, dataset_id, user_id):
    """Create a new model version"""
    try:
        # Create the model version
        response = rekognition_client.create_project_version(
            ProjectArn=project_arn,
            VersionName=f"model_{int(time.time())}",
            OutputConfig={
                'S3Bucket': os.environ.get('S3_BUCKET'),
                'S3KeyPrefix': 'models'
            },
            ManifestS3Bucket=os.environ.get('S3_BUCKET'),
            TrainingData={
                'Assets': [{
                    'GroundTruthManifest': {
                        'S3Object': {
                            'Bucket': os.environ.get('S3_BUCKET'),
                            'Name': dataset_arn.split('/')[-1]
                        }
                    }
                }]
            },
            TestingData={
                'AutoCreate': True
            }
        )
        
        project_version_arn = response['ProjectVersionArn']
        logger.info(f"Created model version: {project_version_arn}")
        
        # Save model to Supabase
        model_id = save_model_to_supabase(project_version_arn, dataset_id, model_name, user_id)
        
        if model_id:
            # Update dataset with model ID
            update_supabase_dataset(dataset_id, 'COMPLETE', model_id)
            
            return {
                'success': True,
                'model_id': model_id,
                'model_arn': project_version_arn
            }
        else:
            logger.error("Failed to save model to Supabase")
            update_supabase_dataset(dataset_id, 'ERROR', None, "Failed to save model to Supabase")
            return {
                'success': False,
                'message': 'Failed to save model to Supabase'
            }
            
    except Exception as e:
        error_message = f"Error creating model: {str(e)}"
        logger.error(error_message)
        update_supabase_dataset(dataset_id, 'ERROR', None, error_message)
        return {
            'success': False,
            'message': error_message
        }

def lambda_handler(event, context):
    """Lambda handler for creating a model from a dataset"""
    try:
        logger.info(f"Starting create model lambda: {json.dumps(event, default=str)}")
        
        # Extract information from event
        dataset_id = event.get('dataset_id')
        dataset_arn = event.get('dataset_arn')
        model_name = event.get('model_name') or f"Model_{int(time.time())}"
        project_name = event.get('project_name')
        user_id = event.get('user_id')
        
        # Do validation
        if not dataset_id or not dataset_arn:
            logger.error("Missing required parameters: dataset_id and dataset_arn")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Missing required parameters',
                    'required': ['dataset_id', 'dataset_arn']
                })
            }
            
        # Update dataset status to PROCESSING
        update_success = update_supabase_dataset(dataset_id, 'PROCESSING')
        if not update_success:
            logger.warning(f"Failed to update dataset {dataset_id} status to PROCESSING")
            
        # Get PROJECT_NAME from environment variables with fallbacks
        if not project_name:
            project_name = os.environ.get('PROJECT_NAME')
            
        # Try to get from SSM if not in environment
        if not project_name:
            project_name = get_ssm_parameter('/datafy-rekognition-stack/project-name')
            
        if not project_name:
            logger.error("No project name found in event, environment, or SSM")
            update_supabase_dataset(dataset_id, 'ERROR', None, "No project name found")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'No project name found'
                })
            }
            
        # Construct the project ARN from the provided name
        aws_region = os.environ.get('AWS_REGION') or 'eu-west-1'
        aws_account_id = os.environ.get('AWS_ACCOUNT_ID') or '587594388832'
        project_arn = f"arn:aws:rekognition:{aws_region}:{aws_account_id}:project/{project_name}"
        
        # Fix project ARN format using the function
        project_arn = fix_project_arn(project_arn)
        if not project_arn:
            logger.error(f"Failed to fix project ARN format for: {project_name}")
            update_supabase_dataset(dataset_id, 'ERROR', None, "Invalid project ARN format")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'message': 'Invalid project ARN format'
                })
            }
            
        logger.info(f"Using project ARN: {project_arn}")
        
        # Create the model
        create_result = create_model(project_arn, dataset_arn, model_name, dataset_id, user_id)
        
        if create_result.get('success'):
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'Model creation started',
                    'model_id': create_result.get('model_id'),
                    'model_arn': create_result.get('model_arn')
                })
            }
        else:
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'message': create_result.get('message')
                })
            }
            
    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}", exc_info=True)
        
        # Try to update dataset status if we have a dataset ID
        if event.get('dataset_id'):
            update_supabase_dataset(
                event.get('dataset_id'), 
                'ERROR', 
                None, 
                f"Exception in create model: {str(e)}"
            )
            
        return {
            'statusCode': 500,
            'body': json.dumps({
                'message': f'Error creating model: {str(e)}'
            })
        }