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

# Initialize AWS SSM client
ssm = boto3.client('ssm')
def get_parameter(name):
    """Get a parameter from AWS SSM Parameter Store"""
    response = ssm.get_parameter(
        Name=name,
        WithDecryption=True
    )
    return response['Parameter']['Value']
def lambda_handler(event, context):
    # Initialize Supabase client 
    url: str = get_parameter('/supabase/url')
    key: str = get_parameter('/supabase/anon')

    # url: str = os.environ.get("SUPABASE_URL")
    # key: str = os.environ.get("SUPABASE_KEY")
    
    supabase: Client = create_client(url, key)
    supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"})
    offset = 0
    limit = 100
    sjreport_records  = []    
    date = event['week_start_date']   

    while True:   
        sjreport = supabase.rpc("get_sjreport_data_pages",{
                        "target_date": date,
                        "page_offset": offset,
                        "limit_count": limit
                    }).execute();
    
        if sjreport.data is None:
            print("No data returned from sjreport.")
            break
        print(f"Number of records sj read : {len(sjreport.data)} + offset : {offset}")
        sjreport_records.extend(sjreport.data)
        if len(sjreport.data) < limit:
            break  # Exit if fewer records than limit are returned
        
        
        offset += limit


    # Convert sjreport data to DataFrame
    sjreport_df = pd.DataFrame(sjreport_records)

    # Prepare to store the combined records
    insertData = []

    # Iterate through sjreport and callsheet DataFrames to create the combined records
    for index, sjElement in sjreport_df.iterrows():        
        # Create the new combined object with NaN checks and .0 removal
        newRow = {
            'chain': str(sjElement['channel']).replace('.0', '') if pd.notna(sjElement['channel']) else '',
            'store_name': str(sjElement['store_name']).replace('.0', '') if pd.notna(sjElement['store_name']) else '',
            'status': str(sjElement['status']).replace('.0', '') if pd.notna(sjElement['status']) else '',            
            'category': str(sjElement['category']).replace('.0', '') if pd.notna(sjElement['category']) else '',
            'brand': str(sjElement['brand']).replace('.0', '') if pd.notna(sjElement['brand']) else '',
            'product': str(sjElement['product']).replace('.0', '') if pd.notna(sjElement['product']) else '',
            'variant': str(sjElement['variant']).replace('.0', '') if pd.notna(sjElement['variant']) else '',
            'size': str(sjElement['size']).replace('.0', '') if pd.notna(sjElement['size']) else '',
            'sjreport_id': sjElement['id'] if pd.notna(sjElement['id']) else None,  
            'store_code': str(sjElement['store_code']).replace('.0', '') if pd.notna(sjElement['store_code']) and sjElement['store_code'] != '' else '',
            'type': 'task',  
            'name': str(sjElement['issue_for_feedback']).replace('.0', '') if pd.notna(sjElement['issue_for_feedback']) else '',
            'store_soh': sjElement['store_soh'] if pd.notna(sjElement['store_soh']) else 0,  
            'system_dros': sjElement['system_dros'] if pd.notna(sjElement['system_dros']) else 0,  
            'scoring_system': str(sjElement['scoring_system']).replace('.0', '') if pd.notna(sjElement['scoring_system']) else '', 
            'last_sold': sjElement['last_sold'] if pd.notna(sjElement['last_sold']) else None,  
            'week_start_date': sjElement['week_start_date'] if pd.notna(sjElement['week_start_date']) else None,  
            'product_id': str(sjElement['product_id']).replace('.0', '') if pd.notna(sjElement['product_id']) else None,  
            'article': str(sjElement['article']).replace('.0', '') if pd.notna(sjElement['article']) else None,
            'customer_id':str(sjElement['customer_id']).replace('.0', '') if pd.notna(sjElement['customer_id']) else None,
            'dc_soh': str(sjElement['dc_soh']).replace('.0', '') if pd.notna(sjElement['dc_soh']) else None
        }
        insertData.append(newRow)

    batch_size = 1000 # batch size for insert
    # Iterate over the insert_data in chunks of batch_size
    for i in range(0, len(insertData), batch_size):
        batch = insertData[i:i + batch_size]  # Get the current batch
        # print(f'tasks : {batch}')
        # Perform the upsert operation
        supabase.table('task').upsert(batch, on_conflict='chain,store_name,product_id,week_start_date,name,sjreport_id').execute()
        
        print(f"Successfully inserted batch {i // batch_size + 1} with {len(batch)} records.")

    print(f"Number of records sjreport: {len(sjreport_records)}")

    # Update issues counts on store table 
    supabase.rpc('update_store_issue_counts').execute()
    return {
        'statusCode': 200,
        'body': json.dumps({'message': 'Data inserted successfully'})
    }