import gradio as gr
import pandas as pd
import tempfile
import os
from openpyxl import load_workbook
from openpyxl.styles import Alignment

def adjust_excel_formatting(file_path):
    wb = load_workbook(file_path)
    ws = wb.active
    for col in ws.columns:
        max_length = 0
        col_letter = col[0].column_letter
        for cell in col:
            if cell.value:
                max_length = max(max_length, len(str(cell.value)))
            cell.alignment = Alignment(wrap_text=True)
        ws.column_dimensions[col_letter].width = max_length + 2
    wb.save(file_path)

def process_file_a_to_b(input_file):
    try:
        # Read the Excel file, skipping the first row (OVERNIGHT/MORNING header)
        input_df = pd.read_excel(input_file.name, header=1)
        
        # Get the date columns (all columns except the first one which contains model names)
        date_columns = input_df.columns[1:].tolist()
        
        # Melt the dataframe to long format
        df_long = input_df.melt(
            id_vars=[input_df.columns[0]],  # First column (Model names)
            var_name='DATE',
            value_name='CHATTER'
        )
        
        # Clean up the data
        df_long = df_long[df_long['CHATTER'].notna()]  # Remove empty cells
        df_long = df_long[df_long['CHATTER'] != '']    # Remove empty strings
        df_long = df_long[df_long['CHATTER'] != 'OFF'] # Remove 'OFF' entries
        
        # Group by chatter and date, collect all models
        grouped = df_long.groupby(['CHATTER', 'DATE'])[input_df.columns[0]].apply(
            lambda x: ', '.join(sorted(x))
        ).reset_index()
        
        # Pivot to get chatters as rows and dates as columns
        pivoted = grouped.pivot(
            index='CHATTER',
            columns='DATE',
            values=input_df.columns[0]
        )
        
        # Reorder columns to match original date order
        pivoted = pivoted[date_columns]
        
        # Define the expected order of chatters
        expected_chatters = [
            'VELJKO2', 'VELJKO3', 'MARKO', 'GODDARD', 'ALEKSANDER', 'FEELIP', 
            'DENIS', 'TOME', 'MILA', 'VELJKO', 'DAMJAN', 'DULE', 'CONRAD', 
            'ALEXANDER', 'VEJKO3'
        ]
        
        # Reindex with expected order, keeping any additional chatters at the end
        final_df = pivoted.reindex(expected_chatters + [x for x in pivoted.index if x not in expected_chatters])
        
        # Fill empty cells with 'OFF'
        final_df = final_df.fillna('OFF')
        
        # Reset index and rename the index column to 'CHATTER'
        final_df = final_df.reset_index()
        final_df = final_df.rename(columns={'index': 'CHATTER'})
        
        return final_df
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})

def process_file_b_to_a(input_file):
    try:
        # Read the Excel file
        input_df = pd.read_excel(input_file.name, header=0)
        
        # Get the date columns (all columns except the first one which contains chatter names)
        date_columns = input_df.columns[1:].tolist()
        
        # Melt the dataframe to long format
        df_long = input_df.melt(
            id_vars=[input_df.columns[0]],  # First column (Chatter names)
            var_name='DATE',
            value_name='MODEL'
        )
        
        # Clean up the data
        df_long = df_long[df_long['MODEL'].notna()]  # Remove empty cells
        df_long = df_long[df_long['MODEL'] != '']    # Remove empty strings
        df_long = df_long[df_long['MODEL'] != 'OFF'] # Remove 'OFF' entries
        
        # Split comma-separated models into separate rows
        df_long['MODEL'] = df_long['MODEL'].str.split(', ')
        df_long = df_long.explode('MODEL')
        
        # Group by model and date, collect all chatters
        grouped = df_long.groupby(['MODEL', 'DATE'])[input_df.columns[0]].apply(
            lambda x: ', '.join(sorted(x))
        ).reset_index()
        
        # Pivot to get models as rows and dates as columns
        pivoted = grouped.pivot(
            index='MODEL',
            columns='DATE',
            values=input_df.columns[0]
        )
        
        # Reorder columns to match original date order
        pivoted = pivoted[date_columns]
        
        # Sort models alphabetically
        final_df = pivoted.sort_index()
        
        # Fill empty cells with 'OFF'
        final_df = final_df.fillna('OFF')
        
        # Reset index to make MODEL a column
        final_df = final_df.reset_index()
        
        return final_df
    except Exception as e:
        return pd.DataFrame({"Error": [str(e)]})

def convert_schedule(file, direction):
    if direction == "Format A → Format B":
        df = process_file_a_to_b(file)
    else:
        df = process_file_b_to_a(file)
    # Save to temp file for download
    with tempfile.NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
        df.to_excel(tmp.name, index=False)
        adjust_excel_formatting(tmp.name)
        tmp.seek(0)
        data = tmp.read()
    return df, (tmp.name,)

def download_file(file_tuple):
    return file_tuple[0]

demo = gr.Blocks()
with demo:
    gr.Markdown("# 📅 Schedule Converter")
    gr.Markdown("Upload your schedule Excel file, select conversion direction, and download the result.")
    with gr.Row():
        file = gr.File(label="Upload Schedule File", type="file")
        direction = gr.Dropdown([
            "Format A → Format B",
            "Format B → Format A"
        ], value="Format A → Format B", label="Conversion Direction")
    with gr.Row():
        process_btn = gr.Button("Process File", variant="primary")
        reset_btn = gr.Button("Upload New File")
    output_table = gr.Dataframe(label="Preview", wrap=True)
    download_button = gr.Button("Download Processed File", visible=False)
    temp_file_path = gr.State(value=None)
    def reset_components():
        return [None, pd.DataFrame(), None, gr.update(visible=False)]
    def process_and_show(file, direction):
        df, out_path = convert_schedule(file, direction)
        if out_path:
            return df, out_path, gr.update(visible=True)
        return df, None, gr.update(visible=False)
    process_btn.click(
        process_and_show,
        inputs=[file, direction],
        outputs=[output_table, temp_file_path, download_button]
    )
    reset_btn.click(
        reset_components,
        outputs=[file, output_table, temp_file_path, download_button]
    )
    download_button.click(
        download_file,
        inputs=temp_file_path,
        outputs=gr.File(label="Processed Schedule")
    )

demo.launch()