# Configuration
#SERVICE_ACCOUNT_FILE = 'digital-pagoda-477703-p7-37d35e143508.json'
#SCOPES = ['https://www.googleapis.com/auth/drive.readonly']



import json
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload

import gradio as gr
import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import PatternFill, Font
from datetime import datetime, timedelta
import io
import tempfile
import os

# ── Utility: Normalize and parse time strings
def fix_time_format(val):
    if pd.isna(val):
        return None
    s = str(val).strip()
    s = s.zfill(4)                # '800' -> '0800'
    return f"{s[:2]}:{s[2:]}"  # '0800' -> '08:00'

def generate_schedule(clients_df, techs_df):
    """Generate schedule based on clients and technicians data"""
    
    # Clean column names
    clients_df.columns = clients_df.columns.str.strip().str.title()
    techs_df.columns = techs_df.columns.str.strip().str.title()

    # Apply time fix and parse datetimes
    for df, cols in [(clients_df, ["Start","End","Nap Time"]), (techs_df, ["Start","End"])]:
        for c in cols:
            if c in df.columns:
                df[c] = df[c].apply(fix_time_format)
                df[c] = pd.to_datetime(df[c], format="%H:%M", errors="coerce")

    # Determine day range
    day_start = min(clients_df["Start"].min(), techs_df["Start"].min())
    day_end = max(clients_df["End"].max(), techs_df["End"].max())
    hours = []
    t = day_start.replace(minute=0)
    while t < day_end:
        hours.append(t)
        t += timedelta(hours=1)

    # Prepare tracking
    schedule = {c: {} for c in clients_df["Name"]}
    tech_busy = {t: set() for t in techs_df["Name"]}
    tech_client_hours = pd.DataFrame(0, index=techs_df["Name"], columns=clients_df["Name"])

    # Assign per hour
    for i, hour in enumerate(hours):
        for _, c in clients_df.iterrows():
            cname = c["Name"]
            if pd.isna(c["Start"]) or pd.isna(c["End"]) or not (c["Start"] <= hour < c["End"]):
                continue

            # Nap time
            nt = c.get("Nap Time")
            if pd.notna(nt) and nt <= hour < nt + timedelta(hours=1):
                schedule[cname][hour] = "NAP"
                continue

            if hour in schedule[cname]:
                continue

            # Build eligible techs
            candidates = []
            for _, t in techs_df.iterrows():
                tname = t["Name"]

                # Check availability
                if pd.isna(t["Start"]) or pd.isna(t["End"]) or not (t["Start"] <= hour < t["End"]):
                    continue
                if c["Gender"] == "F" and t["Gender"] == "M":
                    continue
                if hour in tech_busy[tname]:
                    continue
                if tech_client_hours.at[tname, cname] >= 4:
                    continue

                # Look ahead to find consecutive availability (up to 4 hours max with this client)
                consec = 0
                for j in range(i, min(i+4, len(hours))):
                    next_hour = hours[j]
                    if not (c["Start"] <= next_hour < c["End"]):
                        break
                    if pd.notna(nt) and nt <= next_hour < nt + timedelta(hours=1):
                        break
                    if next_hour >= t["End"] or next_hour in tech_busy[tname]:
                        break
                    if tech_client_hours.at[tname, cname] + consec >= 4:
                        break
                    consec += 1

                # Ignore if tech is only available for 1 hour
                if consec < 2:
                    continue

                prefs = [str(t.get(f"P{i}", "")).strip() for i in (1,2,3,4)]
                if cname in prefs:
                    score = prefs.index(cname) + 1
                elif "Any" in prefs:
                    score = 5
                else:
                    continue
                # Lower score is better; longer block is better → sort by (score, -consec)
                candidates.append((score, -consec, tname))

            if not candidates:
                schedule[cname][hour] = ""
            else:
                candidates.sort()
                _, _, chosen = candidates[0]

                # Re-check how many hours we can assign now (again)
                block = 0
                for j in range(i, min(i+4, len(hours))):
                    next_hour = hours[j]
                    if not (c["Start"] <= next_hour < c["End"]):
                        break
                    if pd.notna(nt) and nt <= next_hour < nt + timedelta(hours=1):
                        break
                    if next_hour >= techs_df.loc[techs_df["Name"] == chosen, "End"].values[0]:
                        break
                    if next_hour in tech_busy[chosen]:
                        break
                    if tech_client_hours.at[chosen, cname] >= 4:
                        break
                    schedule[cname][next_hour] = chosen
                    tech_busy[chosen].add(next_hour)
                    tech_client_hours.at[chosen, cname] += 1
                    block += 1

    return hours, schedule, tech_client_hours

def create_excel_file(hours, schedule, techs_df):
    """Create Excel file with schedule"""
    wb = Workbook()
    ws = wb.active
    ws.title = "Schedule"

    # Header row
    ws.cell(1,1,"Time")
    clients = list(schedule.keys())
    for col, name in enumerate(clients, start=2):
        cell = ws.cell(1, col, name)
        cell.font = Font(bold=True)
        cell.fill = PatternFill("solid", fgColor="C6EFCE")

    # Tech color mapping
    tech_colors = {}
    if "Colours" in techs_df.columns:
        for _, r in techs_df.iterrows():
            col = str(r["Colours"]).strip().lstrip("#").upper()
            if len(col) == 6:
                col = "FF" + col
            tech_colors[r["Name"]] = col

    # Nap style
    nap_fill = PatternFill("solid", fgColor="D9D9D9")

    # Fill schedule
    for row, hour in enumerate(hours, start=2):
        ws.cell(row,1, hour.strftime("%H:%M"))
        for col, cname in enumerate(clients, start=2):
            val = schedule[cname].get(hour, "")
            cell = ws.cell(row, col, val)
            if val == "NAP":
                cell.fill = nap_fill
            elif val in tech_colors:
                fg = tech_colors[val]
                cell.fill = PatternFill("solid", fgColor=fg)

    # Create a temporary file
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".xlsx")
    wb.save(temp_file.name)
    
    return temp_file.name

def schedule_to_dataframe(hours, schedule):
    """Convert schedule to pandas DataFrame for display with Time as first column"""
    # Create DataFrame with time as a regular column
    time_strings = [hour.strftime("%H:%M") for hour in hours]
    
    # Create dictionary for DataFrame
    data = {"Time": time_strings}
    
    # Add each client as a column
    for client_name in schedule.keys():
        client_schedule = []
        for hour in hours:
            value = schedule[client_name].get(hour, "")
            client_schedule.append(value)
        data[client_name] = client_schedule
    
    df = pd.DataFrame(data)
    return df

def process_schedule(techs_file, clients_file):
    """Main function to process the schedule"""
    
    # Check if files are uploaded
    if techs_file is None or clients_file is None:
        return "❌ Please upload both files: 'techs_new_01.csv' and 'clients_01.csv'", None, None
    
    try:
        # Read CSV files
        techs_df = pd.read_csv(techs_file)
        clients_df = pd.read_csv(clients_file)
        
        # Validate required columns
        required_client_cols = ["Name", "Start", "End", "Gender"]
        required_tech_cols = ["Name", "Start", "End", "Gender"]
        
        missing_client = [col for col in required_client_cols if col not in clients_df.columns]
        missing_tech = [col for col in required_tech_cols if col not in techs_df.columns]
        
        if missing_client:
            return f"❌ Missing columns in clients file: {', '.join(missing_client)}", None, None
        if missing_tech:
            return f"❌ Missing columns in techs file: {', '.join(missing_tech)}", None, None
        
        # Generate schedule
        hours, schedule, tech_client_hours = generate_schedule(clients_df, techs_df)
        
        # Create Excel file
        excel_file_path = create_excel_file(hours, schedule, techs_df)
        
        # Convert schedule to DataFrame for display
        schedule_df = schedule_to_dataframe(hours, schedule)
        
        # Create summary
        summary = f"✅ Schedule Generated! Clients: {len(clients_df)}, Techs: {len(techs_df)}, Time: {hours[0].strftime('%H:%M')}-{hours[-1].strftime('%H:%M')}\n📥 Click the arrow → to download Excel file"
        
        return summary, schedule_df, excel_file_path
        
    except Exception as e:
        return f"❌ Error: {str(e)}", None, None

def clear_all():
    """Clear all inputs and outputs"""
    return (
        None,  # techs_file
        None,  # clients_file
        "✅ Ready to upload files",  # output_text
        None,  # schedule_table
        None   # download_file
    )

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🗓️ AI Schedule Generator")
    gr.Markdown("Upload CSV files to generate schedule")
    
    with gr.Row():
        # File uploads side by side
        #gr.Column(scale=0.05)  # Left spacer
        with gr.Column(scale=0.4):
            
            techs_file = gr.File(
                label="Techs.csv",
                file_types=[".csv"],
                type="filepath",
                height=120
                
            )
        with gr.Column(scale=0.4):
            clients_file = gr.File(
                label="Clients.csv", 
                file_types=[".csv"],
                type="filepath",
                height=120
                
            )
        #gr.Column(scale=0.05)  # Right spacer
        
    with gr.Row():
        # Buttons side by side
        with gr.Column(scale=0.4):
            generate_btn = gr.Button("🚀 Generate Schedule", variant="primary", size="lg")
    with gr.Row():
        with gr.Column(scale=0.4):
            refresh_btn = gr.Button("🔄 Clear All", variant="secondary", size="lg")
    
    # Status box - 2 rows maximum
    output_text = gr.Textbox(
        label="Status",
        lines=2,
        max_lines=2,
        interactive=False,
        show_copy_button=True
    )
    
    # Download file component with arrow - this works reliably
    download_file = gr.File(
        label="Download Excel Schedule",
        file_types=[".xlsx"],
        visible=False,
        interactive=True
    )
    
    # Schedule preview
    gr.Markdown("### 📋 AI Schedule Preview")
    schedule_table = gr.Dataframe(
        label="",
        headers=None,
        max_height=400,
        wrap=True,
        interactive=False
    )
    
    # Set up event handlers
    generate_btn.click(
        fn=process_schedule,
        inputs=[techs_file, clients_file],
        outputs=[output_text, schedule_table, download_file]
    ).then(
        # Make download file visible after generation
        lambda: gr.update(visible=True),
        outputs=[download_file]
    )
    
    refresh_btn.click(
        fn=clear_all,
        outputs=[techs_file, clients_file, output_text, schedule_table, download_file]
    ).then(
        # Hide download file after clear
        lambda: gr.update(visible=False),
        outputs=[download_file]
    )
    
    gr.Markdown("""
    ### 📝 Required CSV Formats:
    
    **clients_01.csv:**
    - `Name`, `Start`, `End`, `Gender`, `Nap Time` (optional)
    
    **techs_new_01.csv:**
    - `Name`, `Start`, `End`, `Gender`, `P1-P4` (preferences), `Colours` (optional)
    
    ### 💡 How to download:
    After generating the schedule, click the **arrow (→)** next to "Download Excel Schedule" to download the file.
    """)

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )