import os
import json
from supabase import create_client, Client
import boto3
import pandas as pd
from io import BytesIO
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import logging
from email_service import email_service
import traceback

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Initialize AWS SSM client
ssm = boto3.client('ssm')

def get_parameter(name: str) -> str:
    """Get a parameter from AWS SSM Parameter Store"""
    try:
        response = ssm.get_parameter(
            Name=name,
            WithDecryption=True
        )
        return response['Parameter']['Value']
    except Exception as e:
        logger.error(f"Error getting parameter {name}: {str(e)}")
        raise

def get_supabase_client() -> Client:
    """Initialize and return Supabase client"""
    try:
        url: str = get_parameter('/supabase/url')
        key: str = get_parameter('/supabase/anon')
        supabase: Client = create_client(url, key)
        supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"})
        return supabase
    except Exception as e:
        logger.error(f"Error initializing Supabase client: {str(e)}")
        raise

def get_unprocessed_upload(supabase: Client) -> Optional[Dict]:
    """Get the next unprocessed upload record with times_ran < 5"""
    try:
        upload_response = supabase.table("sjreport_uploads").select("*").eq("processed", "FALSE").lt("times_ran", 5).limit(1).execute()
        return upload_response.data[0] if upload_response.data else None
    except Exception as e:
        logger.error(f"Error getting unprocessed upload: {str(e)}")
        raise

def download_file(supabase: Client, bucket_name: str, file_name: str) -> bytes:
    """Download file from Supabase storage"""
    try:
        response = supabase.storage.from_(bucket_name).download(file_name)
        return response
    except Exception as e:
        logger.error(f"Error downloading file {file_name}: {str(e)}")
        raise

def process_excel_file(file_content: bytes) -> Dict[str, pd.DataFrame]:
    """Process Excel file and return dictionary of DataFrames"""
    try:
        sheets = pd.read_excel(BytesIO(file_content), sheet_name=None, dtype=str)
        for sheet_name in sheets.keys():
            sheets[sheet_name].columns = sheets[sheet_name].columns.str.lower().str.replace(' ', '_')
        return sheets
    except Exception as e:
        logger.error(f"Error processing Excel file: {str(e)}")
        raise

def process_row(row: pd.Series, url: str, sheet_name: str) -> Dict:
    """Process a single row of data"""
    try:
        row_data = row.to_dict()
        modified_row_data = {}

        # Handle null values
        for col, value in row_data.items():
            modified_row_data[col] = "" if pd.isna(value) else value

        # Process SPAR store codes
        if "SPAR" in modified_row_data.get("channel", ""):
            modified_row_data["store_code"] = "S" + modified_row_data["store_code"]

        # Process date
        date_string = modified_row_data["date"][:10]
        date = datetime.strptime(date_string, '%Y-%m-%d').date()
        week_start_date = date + timedelta(days=(7 - date.weekday()))

        # Add metadata
        modified_row_data.update({
            'row': row.name + 1,
            'week_start_date': week_start_date.strftime('%Y-%m-%d'),
            'file': url,
            'sheetname': sheet_name
        })

        return modified_row_data
    except Exception as e:
        logger.error(f"Error processing row: {str(e)}")
        raise

def insert_batch(supabase: Client, table_name: str, batch: List[Dict]) -> None:
    """Insert a batch of records into the database"""
    try:
        response = supabase.table(table_name).upsert(batch, on_conflict='row, week_start_date, file, sheetname').execute()
        if hasattr(response, 'error') and response.error:
            raise Exception(response.error.message)
        logger.info(f"Successfully inserted {len(batch)} records")
    except Exception as e:
        logger.error(f"Error inserting batch: {str(e)}")
        raise

def increment_times_ran(supabase: Client, record_id: str) -> None:
    try:
        # Fetch current value
        response = supabase.table("sjreport_uploads").select("times_ran").eq("id", record_id).execute()
        if response.data and len(response.data) > 0:
            current = response.data[0].get("times_ran", 0) or 0
            new_value = current + 1
            supabase.table("sjreport_uploads").update({"times_ran": new_value}).eq("id", record_id).execute()
    except Exception as e:
        logger.error(f"Error incrementing times_ran: {str(e)}")
        raise

def update_upload_status(supabase: Client, record_id: str, error: str = '', processed: bool = True) -> None:
    """Update the upload record status"""
    try:
        update_data = {
            "error": error,
            "processed": processed
        }
        supabase.table("sjreport_uploads").update(update_data).eq("id", record_id).execute()
    except Exception as e:
        logger.error(f"Error updating upload status: {str(e)}")
        raise

def invoke_lambda_function(lambda_client: boto3.client, week_start_date: str, record_id: str) -> None:
    """Invoke the create tasks Lambda function"""
    try:
        payload = {
            "week_start_date": week_start_date,
            "record_id": record_id
        }
        lambda_client.invoke(
            FunctionName='datafyNew-BatchUpdateSjreport-utMCtdWkdNo8',
            InvocationType='Event',
            Payload=json.dumps(payload)
        )
    except Exception as e:
        logger.error(f"Error invoking create tasks: {str(e)}")
        raise

def lambda_handler(event, context):
    """Main Lambda handler function"""
    try:
        # Initialize clients
        lambda_client = boto3.client('lambda')
        supabase = get_supabase_client()

        # Get unprocessed upload
        record = get_unprocessed_upload(supabase)
        if not record:
            return {
                'statusCode': 400,
                'body': json.dumps({'error': 'No unprocessed uploads found'})
            }

        record_id = record['id']
        bucket_name = "public/" + record['payload']['bucket_name']
        file_name = record['payload']['file_name']
        table_name = record['payload']['path_tokens'][0]

        # Increment times_ran
        increment_times_ran(supabase, record_id)

        # Download and process file
        file_content = download_file(supabase, bucket_name, file_name)
        sheets = process_excel_file(file_content)

        # Process each sheet
        last_week_start_date = None
        BATCH_SIZE = 500  # Reduced batch size for better memory management

        for sheet_name, sheet_df in sheets.items():
            logger.info(f"Processing sheet: {sheet_name}")
            data_to_insert = []

            for index, row in sheet_df.iterrows():
                # Convert each row to a dictionary and add additional fields
                row_data = row.to_dict()
                modified_row_data = {}

                for col, value in row_data.items():
                    if pd.isna(value):
                        modified_row_data[col] = ""
                    else:
                        modified_row_data[col] = value

                # replace row_data with modified_row_data
                row_data = modified_row_data

                date_string = row_data["date"][:10]
                if not date_string or date_string.strip() == "":
                    continue  # Skip this row if date is empty or only whitespace

                try:
                    # Check if "SPAR" is included in "channel" and modify "store_code"
                    if "SPAR" in modified_row_data["channel"]:
                        modified_row_data["store_code"] = "S" + modified_row_data["store_code"]

                    date = datetime.strptime(date_string, '%Y-%m-%d').date()
                    week_start_date = date + timedelta(days=(7 - date.weekday()))
                    last_week_start_date = week_start_date.strftime('%Y-%m-%d')

                    row_data['row'] = index + 1
                    row_data['week_start_date'] = week_start_date.strftime('%Y-%m-%d')
                    row_data['file'] = file_name
                    row_data['sheetname'] = sheet_name

                    data_to_insert.append(row_data)
                except Exception as e:
                    error_details = f"Invalid date format or row error: {str(e)} | Row: {row_data}"
                    logger.error(error_details)
                    # Send error email for row processing error
                    email_service.send_notification(
                        date=datetime.now().strftime('%Y-%m-%d'),
                        status='error',
                        error_message=f"Error processing row:\n{error_details}\n\nStack Trace:\n{traceback.format_exc()}",
                        client_email=False,
                        process_name='Read Files Row Processing'
                    )
                    update_upload_status(supabase, record_id, error_details)
                    continue

                # Insert in smaller batches
                try:
                    if len(data_to_insert) >= BATCH_SIZE:
                        insert_batch(supabase, table_name, data_to_insert)
                        data_to_insert = []
                except Exception as e:
                    error_details = f"Error inserting data: {str(e)} | Batch: {data_to_insert}"
                    logger.error(error_details)
                    # Send error email for batch insertion error
                    # User-friendly error message for client
                    user_friendly_message = "An error occurred during the SJREPORT import."
                    if "Could not find the" in str(e) and "column" in str(e):
                        import re
                        match = re.search(r"Could not find the '([^']+)' column", str(e))
                        if match:
                            missing_col = match.group(1).replace('_', ' ').upper()
                            user_friendly_message = f'Could not find column "{missing_col}" in the database. Please fix your file and try again.'
                    
                    # Add file name to the user-friendly message
                    file_info = f'File: {file_name}\n'
                    user_friendly_message = file_info + user_friendly_message
                    
                    email_service.send_notification(
                        date=datetime.now().strftime('%Y-%m-%d'),
                        status='error',
                        error_message=user_friendly_message,
                        client_email=True,
                        process_name='SJREPORT Import'
                    )
                    update_upload_status(supabase, record_id, error_details)
                    return {
                        'statusCode': 400,
                        'body': json.dumps({'error': error_details})
                    }

            # Insert remaining records
            if data_to_insert:
                insert_batch(supabase, table_name, data_to_insert)
                logger.info(f"Inserted {len(data_to_insert)} records into the table.")

        # Update product table and create tasks
        supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"})
        # Ignore duplicate-key violations when products already exist
        try:
            supabase.rpc("insert_distinct_products").execute()
        except Exception as e:
            if "23505" in str(e) or "duplicate key value violates unique constraint" in str(e):
                logger.warning("insert_distinct_products: duplicate products detected; continuing without failing the run")
            else:
                raise
        supabase.rpc("update_product_customer_ids").execute()
        if last_week_start_date:
            invoke_lambda_function(lambda_client, last_week_start_date, record_id)

        # Mark upload as processed
        update_upload_status(supabase, record_id, processed=True)

        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'Data processed successfully'})
        }

    except Exception as e:
        logger.error(f"Error in lambda_handler: {str(e)}")
        error_message = f"Error Type: {type(e).__name__}\nError Message: {str(e)}\n\nStack Trace:\n{traceback.format_exc()}"
        
        # Send error email to admin
        email_service.send_notification(
            date=datetime.now().strftime('%Y-%m-%d'),
            status='error',
            error_message=error_message,
            client_email=False,
            process_name='Read Files Processing'
        )
        
        # Send error email to client
        user_friendly_message = f"An error occurred during the SJREPORT import.\nFile: {file_name}\nError: {str(e)}"
        email_service.send_notification(
            date=datetime.now().strftime('%Y-%m-%d'),
            status='error',
            error_message=user_friendly_message,
            client_email=True,
            process_name='SJREPORT Import'
        )
        
        # If it's a lock timeout
        if hasattr(e, 'args') and e.args and '55P03' in str(e.args[0]):
            update_upload_status(supabase, record_id, f"Error in lambda_handler: {str(e)}")
        else:
            update_upload_status(supabase, record_id, f"Error in lambda_handler: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }