import pandas as pd
import openpyxl  # ensure XLSX engine is available
import gradio as gr
import uuid
import os

# 7-Day Schedule Converter with explicit iteration and header fill and download support

def convert_schedule(file_path, direction):
    # 1. Load raw header rows to determine day labels
    raw = pd.read_excel(file_path, header=None)
    header1 = raw.iloc[0, 1:].astype(object)
    header2 = raw.iloc[1, 1:].astype(object)

    # Decide which header row to use: prefer second if fully populated
    if header2.notna().all() and not header2.str.startswith('Unnamed').any():
        days = header2.tolist()
        data_start = 2
    else:
        # Forward-fill merged first-row headers
        days = []
        last = None
        for val in header1:
            if pd.isna(val) or str(val).startswith('Unnamed'):
                days.append(last)
            else:
                last = str(val)
                days.append(last)
        data_start = 1

    # 2. Load actual data using resolved day columns
    df = pd.read_excel(
        file_path,
        header=data_start,
        index_col=0,
        usecols=[0] + list(range(1, len(days) + 1))
    )
    df.columns = [str(day) for day in days]

    # 3. Retain original day column order
    day_cols = list(df.columns)

    # 4. Build assignment mapping via explicit iteration
    assignments = {}
    if direction == 'A to B':
        # Models in rows → Texters as rows
        for model in df.index.astype(str):
            for day in day_cols:
                cell = df.at[model, day]
                for texter in str(cell).split(','):
                    texter = texter.strip()
                    if not texter or texter.lower() in ['nan', 'none']:
                        continue
                    assignments.setdefault(texter, {d: [] for d in day_cols})
                    assignments[texter][day].append(model)
        index = sorted(assignments.keys())
        result = pd.DataFrame(index=index, columns=day_cols)
        first_col_name = 'Texter'
        for texter, days_map in assignments.items():
            for day in day_cols:
                models = days_map.get(day, [])
                result.at[texter, day] = ', '.join(models) if models else 'OFF'
    else:
        # Texters in rows → Models as rows
        for texter in df.index.astype(str):
            for day in day_cols:
                cell = df.at[texter, day]
                for model in str(cell).split(','):
                    model = model.strip()
                    if not model or model.lower() in ['nan', 'none']:
                        continue
                    assignments.setdefault(model, {d: [] for d in day_cols})
                    assignments[model][day].append(texter)
        index = sorted(assignments.keys())
        result = pd.DataFrame(index=index, columns=day_cols)
        first_col_name = 'Model'
        for model, days_map in assignments.items():
            for day in day_cols:
                texters = days_map.get(day, [])
                result.at[model, day] = ', '.join(texters) if texters else 'OFF'

    # 5. Cleanup axis names
    result.index.name = None
    result.columns.name = None

    # For display, include index as a column
    display_df = result.reset_index().rename(columns={'index': first_col_name})

    # 6. Save to Excel for download (relative path so Gradio can serve it)
    output_filename = f"converted_{uuid.uuid4().hex}.xlsx"
    result.to_excel(output_filename, engine='openpyxl', index=True)

    # Return both DataFrame and download path
    return display_df, output_filename

# Build Gradio interface and launch immediately for Hugging Face Spaces
iface = gr.Interface(
    fn=convert_schedule,
    inputs=[
        gr.File(label='Upload Weekly Schedule (.xlsx)', file_count='single', type='filepath'),
        gr.Radio(['A to B', 'B to A'], label='Convert Direction')
    ],
    outputs=[
        gr.Dataframe(label='Converted Schedule'),
        gr.File(label='Download Converted Excel')
    ],
    title='7-Day Schedule Converter',
    description=(
        'Upload a 7-column weekly schedule (Models vs Days) with merged or single headers, '
        'then flip between Models→Texters or Texters→Models. Download the result as .xlsx.'
    ),
    flagging_mode='never'
)
# Launch on 0.0.0.0:7860 for Spaces
iface.launch(server_name='0.0.0.0', server_port=7860)