src/formats.py

import pandas as pd
import os
import logging
from datetime import datetime, date

logger = logging.getLogger(__name__)

def calculate_passenger_type(dob_str):
    """
    Calculate passenger type based on age from date of birth.
    
    Age groups:
    - Infants: 0-1 year (0-11 months)
    - Children: 1-17 years
    - Adults: 18+ years
    
    Returns:
        For Fly Dubai: "Infant", "Child", "Adult"
        For Iraqi/Fly Baghdad: "INF", "CHD", "ADT"
    """
    if not dob_str:
        return "Adult", "ADT"  # Default to adult if no DOB
    
    try:
        # Parse the date (assuming DD/MM/YYYY format)
        dob = datetime.strptime(dob_str, '%d/%m/%Y').date()
        today = date.today()
        
        # Calculate age in years
        age_years = today.year - dob.year
        if today.month < dob.month or (today.month == dob.month and today.day < dob.day):
            age_years -= 1
        
        # Calculate age in months for more accurate infant check
        age_months = (today.year - dob.year) * 12 + (today.month - dob.month)
        if today.day < dob.day:
            age_months -= 1
        
        # Determine passenger type based on age
        if age_months < 12:  # Less than 12 months - Infant
            return "Infant", "INF"
        elif age_years < 18:  # 1-17 years - Child
            return "Child", "CHD"
        else:  # 18+ years - Adult
            return "Adult", "ADT"
            
    except (ValueError, TypeError):
        return "Adult", "ADT"  # Default to adult if date parsing fails

def format_fly_dubai(data_list):
    """
    Formats data for Fly Dubai template (previously Iraqi Airways).
    Columns: TYPE, TITLE, FIRST NAME, LAST NAME, DOB (DD/MM/YYYY), GENDER
    Example: Adult MR FirstNameOne LastNameOne 13/8/2015 Male
    """
    formatted_rows = []
    
    for item in data_list:
        sex = item.get('sex', '').upper()
        # Map Gender to Male/Female
        gender_full = "Male" if sex == 'M' else "Female"
        title = "MR" if sex == 'M' else "MRS"
        
        # Calculate passenger type based on age
        raw_dob = item.get('date_of_birth', '')
        
        # Convert MRZ date (YYMMDD) to DD/MM/YYYY for age calculation
        ddmmyyyy_dob = ''
        if raw_dob and len(raw_dob) == 6:  # YYMMDD format
            try:
                from datetime import datetime
                mrz_date = datetime.strptime(raw_dob, '%y%m%d').date()
                ddmmyyyy_dob = mrz_date.strftime('%d/%m/%Y')
            except:
                ddmmyyyy_dob = raw_dob
        else:
            ddmmyyyy_dob = raw_dob
            
        passenger_type, _ = calculate_passenger_type(ddmmyyyy_dob)
        
        row = {
            "TYPE": passenger_type,
            "TITLE": title,
            "FIRST NAME": item.get('name', ''),
            "LAST NAME": item.get('surname', ''),
            "DOB (DD/MM/YYYY)": ddmmyyyy_dob,
            "GENDER": gender_full
        }
        formatted_rows.append(row)
        
    return pd.DataFrame(formatted_rows)

from datetime import datetime

def _to_ddmmmyy(date_str):
    """Converts DD/MM/YYYY to DDMMMYY format, e.g., 13NOV84."""
    if not date_str:
        return ""
    try:
        # Parse from DD/MM/YYYY
        dt_obj = datetime.strptime(date_str, '%d/%m/%Y')
        # Format to DDMMMYY with strict zero padding for day
        # %d should do it, but to be safe we can use string formatting
        day = dt_obj.strftime('%d')
        month_year = dt_obj.strftime('%b%y').upper()
        return f"{day}{month_year}"
    except ValueError:
        return date_str # Return original if parsing fails

def format_iraqi(data_list):
    """
    Formats data for Iraqi template (previously Flydubai).
    - Date of Birth: DDMMMYY (e.g., 13NOV84)
    - Passport Expiry Date: DDMMMYY (e.g., 13NOV84)
    """
    formatted_rows = []
    
    for item in data_list:
        sex = item.get('sex', '').upper()
        title = "MR" if sex == 'M' else "MRS" 
        
        first_middle = item.get('name', '')
        
        # Get date data for age calculation
        raw_dob = item.get('date_of_birth', '')
        raw_expiry = item.get('expiration_date', '')
        
        # Calculate passenger type based on age
        # The date might be in DDMMMYY format (18NOV16) or YYMMDD format (161118)
        ddmmyyyy_dob = ''
        if raw_dob:
            try:
                # First try YYMMDD format (MRZ raw format) - 6 chars
                if len(raw_dob) == 6:
                    mrz_date = datetime.strptime(raw_dob, '%y%m%d').date()
                    ddmmyyyy_dob = mrz_date.strftime('%d/%m/%Y')
                else:
                    # Try DDMMMYY format (already formatted) - 7+ chars
                    formatted_date = datetime.strptime(raw_dob, '%d%b%y').date()
                    ddmmyyyy_dob = formatted_date.strftime('%d/%m/%Y')
            except:
                ddmmyyyy_dob = raw_dob
        else:
            ddmmyyyy_dob = raw_dob
            
        _, ptc_code = calculate_passenger_type(ddmmyyyy_dob)
        
        # Use the dates as-is since they're already formatted
        dob_formatted = raw_dob
        expiry_formatted = raw_expiry
        
        row = {
            "Last Name": item.get('surname', ''),
            "First Name and Middle Name": first_middle,
            "Title": title,
            "PTC": ptc_code, # Passenger Type Code based on age
            "Gender": sex,
            "Date of Birth": dob_formatted,
            "Passport Last Name": item.get('surname', ''),
            "Passport First Name": first_middle,
            "Passport Middle Name": "",
            "Passport Number": item.get('passport_number', ''),
            "Passport Nationality": item.get('nationality', '')[:3],
            "Passport Issue Country": item.get('country', '')[:3],
            "Passport Expiry Date": expiry_formatted,
            # Empty fields
            "Visa Number": "", "Visa Type": "", "Visa Issue Date": "", "Place of Birth": "",
            "Visa Place of Issue": "", "Visa Country of Application": "", "Address Type": "",
            "Address Country": "", "Address Details": "", "Address City": "",
            "Address State": "", "Address Zip Code": ""
        }
        formatted_rows.append(row)
        
    return pd.DataFrame(formatted_rows)

def format_fly_baghdad(data_list):
    """
    Formats data for Fly Baghdad template.
    Columns: Sequence, Traveling With, Pax Type, Title, First Name, Last Name, Gender, 
             DOB (dd/mm/yyyy), Nationality, Passport Number, Passport Expiry (dd/mm/yyyy), 
             Passport Issued Country
    """
    from config.settings import COUNTRY_CODES
    
    # Create a mapping for quick lookup
    country_map = {c['alpha-3']: c['name'].upper() for c in COUNTRY_CODES}
    
    formatted_rows = []
    
    for i, item in enumerate(data_list):
        sex = item.get('sex', '').upper()
        gender_full = "MALE" if sex == 'M' else "FEMALE"
        title = "MR" if sex == 'M' else "MRS"
        
        # Get and format DOB
        raw_dob = item.get('date_of_birth', '')
        ddmmyyyy_dob = ''
        if raw_dob and len(raw_dob) == 6:
            try:
                mrz_date = datetime.strptime(raw_dob, '%y%m%d').date()
                ddmmyyyy_dob = mrz_date.strftime('%d/%m/%Y')
            except:
                ddmmyyyy_dob = raw_dob
        else:
            ddmmyyyy_dob = raw_dob
            
        # Get and format Expiry
        raw_expiry = item.get('expiration_date', '')
        ddmmyyyy_expiry = ''
        if raw_expiry and len(raw_expiry) == 6:
            try:
                mrz_date = datetime.strptime(raw_expiry, '%y%m%d').date()
                ddmmyyyy_expiry = mrz_date.strftime('%d/%m/%Y')
            except:
                ddmmyyyy_expiry = raw_expiry
        else:
            ddmmyyyy_expiry = raw_expiry
            
        _, ptc_code = calculate_passenger_type(ddmmyyyy_dob)
        
        # Map nationality and country
        nat_code = item.get('nationality', '')[:3].upper()
        country_code = item.get('country', '')[:3].upper()
        
        nationality = country_map.get(nat_code, nat_code)
        issued_country = country_map.get(country_code, country_code)
        
        row = {
            "Sequence": i + 1,
            "Traveling With": "",
            "Pax Type": ptc_code,
            "Title": title,
            "First Name": item.get('name', ''),
            "Last Name": item.get('surname', ''),
            "Gender": gender_full,
            "DOB (dd/mm/yyyy)": ddmmyyyy_dob,
            "Nationality": nationality,
            "Passport Number": item.get('passport_number', ''),
            "Passport Expiry (dd/mm/yyyy)": ddmmyyyy_expiry,
            "Passport Issued Country": issued_country
        }
        formatted_rows.append(row)
        
    return pd.DataFrame(formatted_rows)


def export_to_spreadsheet(data_list, output_file, format='excel'):
    """
    Exports a list of dictionaries to a spreadsheet (Excel or CSV).
    """
    if not data_list:
        logger.warning("No data to export.")
        return False

    try:
        df = pd.DataFrame(data_list)
        
        # Ensure output directory exists
        os.makedirs(os.path.dirname(output_file), exist_ok=True)

        if format.lower() == 'excel' or output_file.endswith('.xlsx'):
            if not output_file.endswith('.xlsx'):
                output_file += '.xlsx'
            df.to_excel(output_file, index=False)
            logger.info(f"Data exported to {output_file}")
            
        elif format.lower() == 'csv' or output_file.endswith('.csv'):
            if not output_file.endswith('.csv'):
                output_file += '.csv'
            df.to_csv(output_file, index=False)
            logger.info(f"Data exported to {output_file}")
            
        else:
            logger.error(f"Unsupported format: {format}")
            return False
            
        return True

    except Exception as e:
        logger.error(f"Failed to export data: {e}")
        return False