import asyncio
import os
import sys
import pandas as pd
from datetime import datetime
import calendar

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.future import select
from dotenv import load_dotenv

load_dotenv()
DATABASE_URL = os.environ.get("DATABASE_URL")
engine = create_async_engine(DATABASE_URL)
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

from app.models.maintenance_models import MaintenanceAsset, MaintenancePMSchedule

DIR_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "Maintainance"))

MONTH_MAP = {
    'JAN': 1, 'FEB': 2, 'MAR': 3, 'APR': 4, 'MAY': 5, 'JUN': 6,
    'JUL': 7, 'AUG': 8, 'SEP': 9, 'OCT': 10, 'NOV': 11, 'DEC': 12
}

# Map exact spellings from excel to DB types
MACHINE_TYPES = {
    'CO 2 M/C': 'CO2_MIG',
    'OVAN': 'OVAN',
    'COMPRESSER': 'COMPRESSOR',
    'BLOWER': 'BLOWER',
    'HOIST': 'HOIST',
    'SHOWER': 'SHOWER'
}

async def get_all_assets(session):
    result = await session.execute(select(MaintenanceAsset))
    assets = result.scalars().all()
    # Create lookup map (type, num) -> id
    lookup = {}
    for a in assets:
        lookup[(a.machine_type, a.machine_number)] = a.id
    return lookup

async def process_pm_schedule():
    path = os.path.join(DIR_PATH, "PM SUDHULE 2026.xlsx")
    print(f"\n--- Processing {path} ---")
    xl = pd.ExcelFile(path)

    async with SessionLocal() as session:
        assets_lookup = await get_all_assets(session)
        schedules_to_add = []

        # Only process 2026 sheets to match our DB constraints and recent history
        sheets_2026 = [s for s in xl.sheet_names if '2026' in s]
        
        for sheet_name in sheets_2026:
            print(f" Reading sheet {sheet_name}...")
            
            # Determine month from sheet name (e.g. 'JAN.2026')
            month_str = ''.join(filter(str.isalpha, sheet_name)).upper()
            month = None
            for k, v in MONTH_MAP.items():
                if k in month_str:
                    month = v
                    break
            
            if not month:
                print(f"  Could not determine month from '{sheet_name}'")
                continue
                
            year = 2026

            df = pd.read_excel(path, sheet_name=sheet_name, header=None)
            
            # Find the header row (contains "MACH NAME" and dates 1-31)
            header_row_idx = -1
            date_col_start = -1
            
            for i, row in df.iterrows():
                vals = [str(v).upper() for v in row.values]
                if 'MACH NAME' in vals:
                    header_row_idx = i
                    date_col_start = vals.index('MACH NAME') + 1
                    break
            
            if header_row_idx == -1:
                print("  Could not find header row.")
                continue

            current_machine_type = None
            
            # Loop through rows after header
            for i in range(header_row_idx + 1, len(df)):
                row_vals = df.iloc[i].values
                col1 = str(row_vals[1]).strip() if len(row_vals) > 1 else ""
                
                # Check if this row declares a new machine block
                for excel_name, db_type in MACHINE_TYPES.items():
                    if excel_name in col1.upper() or excel_name.replace(' ', '') in col1.upper().replace(' ', ''):
                        current_machine_type = db_type
                        break
                
                if not current_machine_type:
                    continue

                # Read dates from date_col_start onwards
                for col_idx in range(date_col_start, len(row_vals)):
                    # Get the date day from the header row
                    day_val = str(df.iloc[header_row_idx, col_idx]).strip()
                    day_num = ''.join(filter(str.isdigit, day_val.split('.')[0])) # Handle '11.0'
                    if not day_num:
                        continue
                        
                    day = int(day_num)
                    if day < 1 or day > 31:
                        continue
                        
                    cell_val = str(row_vals[col_idx]).strip()
                    if cell_val.lower() == 'nan' or not cell_val:
                        continue
                        
                    # Ignore Sunday text
                    if cell_val.upper() in ['S', 'U', 'N', 'D', 'A', 'Y']:
                        continue
                        
                    # Ignore 'CO', 'O', 'BL', 'COM', 'H' text which is just formatting
                    if not any(char.isdigit() for char in cell_val):
                        continue

                    # Extract the machine number
                    machine_num_str = ''.join(filter(str.isdigit, cell_val))
                    if not machine_num_str:
                        continue
                        
                    machine_num = int(machine_num_str)
                    
                    # Look up asset
                    asset_id = assets_lookup.get((current_machine_type, machine_num))
                    if not asset_id:
                        continue
                        
                    try:
                        schedule_date = datetime(year, month, day).date()
                        schedules_to_add.append({
                            "asset_id": asset_id,
                            "scheduled_date": schedule_date,
                            "year": year,
                            "month": month,
                            "status": "PLANNED"
                        })
                    except ValueError:
                        pass # Invalid date like Feb 30

        print(f" Found {len(schedules_to_add)} PM Schedule entries to insert.")
        
        # Deduplicate and insert
        seen = set()
        to_insert = []
        for s in schedules_to_add:
            key = (s["asset_id"], s["scheduled_date"])
            if key not in seen:
                seen.add(key)
                to_insert.append(MaintenancePMSchedule(**s))
                
        session.add_all(to_insert)
        await session.commit()
        print(" Successfully inserted into DB.")

if __name__ == "__main__":
    asyncio.run(process_pm_schedule())