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:
        input_df = pd.read_excel(input_file.name, header=1)
        date_columns = input_df.columns[1:].tolist()
        df_long = input_df.melt(
            id_vars=[input_df.columns[0]],
            var_name='DATE',
            value_name='CHATTER'
        )
        df_long = df_long[df_long['CHATTER'].notna()]
        df_long = df_long[df_long['CHATTER'] != '']
        df_long = df_long[df_long['CHATTER'] != 'OFF']
        grouped = df_long.groupby(['CHATTER', 'DATE'])[input_df.columns[0]].apply(
            lambda x: ', '.join(sorted(x))
        ).reset_index()
        pivoted = grouped.pivot(
            index='CHATTER',
            columns='DATE',
            values=input_df.columns[0]
        )
        pivoted = pivoted[date_columns]
        # Use the order of chatters as they appear in the input file
        input_order = grouped['CHATTER'].drop_duplicates().tolist()
        final_df = pivoted.reindex(input_order)
        final_df = final_df.fillna('OFF')
        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:
        input_df = pd.read_excel(input_file.name, header=0)
        date_columns = input_df.columns[1:].tolist()
        df_long = input_df.melt(
            id_vars=[input_df.columns[0]],
            var_name='DATE',
            value_name='MODEL'
        )
        df_long = df_long[df_long['MODEL'].notna()]
        df_long = df_long[df_long['MODEL'] != '']
        df_long = df_long[df_long['MODEL'] != 'OFF']
        df_long['MODEL'] = df_long['MODEL'].str.split(', ')
        df_long = df_long.explode('MODEL')
        grouped = df_long.groupby(['MODEL', 'DATE'])[input_df.columns[0]].apply(
            lambda x: ', '.join(sorted(x))
        ).reset_index()
        pivoted = grouped.pivot(
            index='MODEL',
            columns='DATE',
            values=input_df.columns[0]
        )
        pivoted = pivoted[date_columns]
        # Use the order of models as they appear in the input file
        input_order = grouped['MODEL'].drop_duplicates().tolist()
        final_df = pivoted.reindex(input_order)
        final_df = final_df.fillna('OFF')
        final_df = final_df.reset_index()
        final_df = final_df.rename(columns={'index': 'MODEL'})
        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="filepath")
        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()