import pandas as pd
from openpyxl import load_workbook
from openpyxl.utils import get_column_letter
import gradio as gr
import os
import warnings

# Suppress openpyxl warnings
warnings.filterwarnings("ignore", category=UserWarning, module="openpyxl")

# Load the constant mapping file (embedded in the app)
def load_mapping():
    mapping_path = "Levi's Data Mapping.xlsx"
    return pd.read_excel(mapping_path)

# Function to extract and map data from the input workbook
def transform_data(input_path, mapping_df):
    # Load the input workbook
    input_workbook = pd.ExcelFile(input_path)

    # Initialize a dictionary to store data for output
    output_data = {}

    # Iterate through each mapping rule
    for _, row in mapping_df.iterrows():
        output_column = row["PO Output Column"]
        input_sheet = row["Sheet Name"]
        input_column = row["Input Column"]
        start_row = row.get("Start Row", 2)  # Default start row is 2 if not specified

        if pd.isna(output_column) or pd.isna(input_sheet) or pd.isna(input_column):
            continue

        # Extract data from the specified sheet and column
        if input_sheet in input_workbook.sheet_names:
            sheet_data = pd.read_excel(input_path, sheet_name=input_sheet, usecols=[input_column], skiprows=start_row - 1)
            output_data[output_column] = sheet_data[input_column].tolist()
        else:
            output_data[output_column] = []  # If sheet is missing, add empty column

    # Ensure all columns have the same number of rows by filling with blanks
    max_rows = max(len(col_data) for col_data in output_data.values())
    for key in output_data:
        while len(output_data[key]) < max_rows:
            output_data[key].append("")

    return pd.DataFrame(output_data)

# Main processing function
def process_files(input_workbook):
    try:
        # Validate input file
        if not os.path.exists(input_workbook):
            return None, "Input workbook file does not exist."

        # Load the constant mapping data
        mapping_df = load_mapping()

        # Transform the data
        transformed_data = transform_data(input_workbook, mapping_df)

        # Load the output template (embedded in the app)
        output_template_path = "Generated_Output.xlsx"
        if not os.path.exists(output_template_path):
            return None, "Output template file is missing."

        output_workbook = load_workbook(output_template_path)
        output_sheet = output_workbook.active

        # Write transformed data to the output sheet
        for row_idx, row_data in enumerate(transformed_data.itertuples(index=False), start=2):
            for col_idx, value in enumerate(row_data, start=1):
                output_sheet.cell(row=row_idx, column=col_idx, value=value)

        # Save the generated output file
        output_file_path = "Generated_Output_Final.xlsx"
        output_workbook.save(output_file_path)

        return output_file_path, None

    except Exception as e:
        return None, f"An error occurred during file generation: {e}"

# Define the Gradio interface
def generate_excel(input_workbook):
    file_path, error = process_files(input_workbook.name)
    if error:
        return None, error
    return file_path, "File generated successfully."

with gr.Blocks() as app:
    gr.Markdown("# Excel Sheet Generator")
    gr.Markdown("Upload the input workbook to generate the final Excel file.")

    with gr.Row():
        input_workbook = gr.File(label="Input Workbook", file_types=[".xlsx"])

    output_file = gr.File(label="Generated Excel File")
    status_message = gr.Textbox(label="Status", interactive=False)

    generate_button = gr.Button("Generate Excel File")

    generate_button.click(generate_excel, inputs=[input_workbook], outputs=[output_file, status_message])

# Launch the app
app.launch()