import gradio as gr
import pandas as pd
import re
import os

# ---------- Helpers ----------

def sanitize_string(s):
    return re.sub(r'[^A-Z0-9]', '', s.upper())

def extract_letters(s, count):
    letters = re.sub(r'[^A-Z]', '', s.upper())
    return letters[:count].ljust(count, '0')

def extract_street_number(text: str) -> str:
    text_upper = text.upper()
    patterns = [
        r'ROOM\s+(\d{1,4})',
        r'BUILDING\s+(\d{1,4})',
        r'NO\.?\s+(\d{1,4})',
    ]
    for pattern in patterns:
        match = re.search(pattern, text_upper)
        if match:
            return match.group(1)
    fallback = re.search(r'\d{1,4}', text_upper)
    return fallback.group(0) if fallback else ''

def extract_city(text: str) -> str:
    city_match = re.search(r'([A-Z]+) CITY', text.upper())
    return city_match.group(1) if city_match else "UNKNOWN"

def extract_company_letters(name: str) -> str:
    words = name.upper().split()
    first = extract_letters(words[0], 3) if len(words) > 0 else ''
    second = extract_letters(words[1], 3) if len(words) > 1 else ''
    return (first + second)[:6].ljust(3, '0')

def generate_mid(country_code, company_name, full_address):
    if not country_code or not company_name or not full_address:
        return None
    cc = sanitize_string(country_code)[:2]
    if len(cc) != 2:
        return None
    company_part = extract_company_letters(company_name)
    street_number = extract_street_number(full_address)
    city = extract_city(full_address)
    city_part = extract_letters(city, 3)
    return f"{cc}{company_part}{street_number}{city_part}"

# ---------- Manual Entry ----------

def manual_mid(full_address, manual_company_name, country_code):
    if not full_address.strip():
        return "❌ Please provide a full address."

    cc = sanitize_string(country_code or "CN")[:2]
    if len(cc) != 2:
        return "❌ Invalid country code."

    company_part = extract_company_letters(manual_company_name or "")
    street_number = extract_street_number(full_address)
    city = extract_city(full_address)
    city_part = extract_letters(city, 3)

    mid = f"{cc}{company_part}{street_number}{city_part}"

    return (
        f"✅ MID: {mid}\n\n"
        f"Company: {manual_company_name or '(auto)'}\n"
        f"Street #: {street_number or '(not found)'}\n"
        f"City: {city}\n"
        f"Country: {cc}"
    )

# ---------- File Upload Handling ----------

def process_excel(file_obj):
    try:
        df = pd.read_excel(file_obj)

        if not all(col in df.columns for col in ["CountryCode", "CompanyName", "FullAddress"]):
            return "❌ Missing required columns: CountryCode, CompanyName, FullAddress", None

        df["MID"] = df.apply(
            lambda row: generate_mid(
                row["CountryCode"],
                row["CompanyName"],
                row["FullAddress"]
            ),
            axis=1
        )

        output_path = "generated_mid_output.xlsx"
        df.to_excel(output_path, index=False)
        return "✅ MID generation completed. Download the file below.", output_path

    except Exception as e:
        return f"❌ Error processing file: {str(e)}", None

# ---------- Gradio UI ----------

with gr.Blocks(title="MID Generator") as demo:
    gr.Markdown("## 🏷️ CBP MID Code Generator")

    with gr.Tab("🔤 Manual Input"):
        full_address = gr.Textbox(label="Full Address", lines=3)
        company_name = gr.Textbox(label="Company Name", placeholder="Optional (auto-detect if blank)")
        country_code = gr.Textbox(label="Country Code", value="CN")

        mid_output = gr.Textbox(label="Generated MID", lines=6)
        gen_btn = gr.Button("Generate MID")

        gen_btn.click(
            fn=manual_mid,
            inputs=[full_address, company_name, country_code],
            outputs=mid_output
        )

    with gr.Tab("📁 Upload Excel File"):
        gr.Markdown("**Upload an Excel file with the following columns:** `CountryCode`, `CompanyName`, `FullAddress`")
        excel_input = gr.File(file_types=[".xlsx"], label="Upload Excel File")
        file_status = gr.Textbox(label="Status", interactive=False)
        download_link = gr.File(label="Download Output")

        excel_input.change(
            fn=process_excel,
            inputs=[excel_input],
            outputs=[file_status, download_link]
        )

    with gr.Tab("📄 Template"):
        gr.Markdown("⬇️ Download the sample Excel template:")
        gr.File(value="mid_upload_template.xlsx", label="Download Template")

demo.launch()