import pandas as pd
from openpyxl import Workbook
import gradio as gr
import os
import warnings

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

# Hardcoded mapping logic (from reference file)
MAPPING = [
    {"Sheet Name": "Sheet1", "Input Column": "A", "PO Output Column": "Order Number", "Start Row": 2},
    {"Sheet Name": "Sheet1", "Input Column": "B", "PO Output Column": "Customer Name", "Start Row": 2},
    {"Sheet Name": "Sheet2", "Input Column": "C", "PO Output Column": "Product Code", "Start Row": 2},
    {"Sheet Name": "Sheet2", "Input Column": "D", "PO Output Column": "Quantity", "Start Row": 2},
    # Add more mappings as needed
]

# Function to extract and map data from the input workbook
def transform_data(input_path):
    # 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 mapping in MAPPING:
        output_column = mapping["PO Output Column"]
        input_sheet = mapping["Sheet Name"]
        input_column = mapping["Input Column"]
        start_row = mapping.get("Start Row", 2)  # Default start row is 2 if not specified

        # 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.iloc[:, 0].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."

        # Transform the data
        transformed_data = transform_data(input_workbook)

        # Create a new workbook for the output
        output_workbook = Workbook()
        output_sheet = output_workbook.active
        output_sheet.title = "Generated Output"

        # Write headers
        headers = transformed_data.columns.tolist()
        for col_idx, header in enumerate(headers, start=1):
            output_sheet.cell(row=1, column=col_idx, value=header)

        # Write data rows
        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()