#!/usr/bin/env python3 """ Entry Milestone Exporter Module for HF Spaces Exports entry milestone data from Notion to Excel format (export-only version). """ import pandas as pd import sys from datetime import datetime from typing import List, Dict, Optional import pytz # Add src directory to path sys.path.append('src') from notion_integration import NotionIntegration from constants import POD_TIMEZONE_MAPPINGS class EntryMilestoneExporter: """Class to export entry milestone data to Excel format.""" @staticmethod def convert_utc_to_pod_timezone(utc_dt: datetime, pod: Optional[str]) -> datetime: """ Convert UTC datetime to POD local timezone without timezone info. Args: utc_dt: UTC datetime object pod: Port of discharge (JFK, LAX, ORD, MIA) Returns: Timezone-naive datetime in POD local time (or UTC if POD not found) """ if not utc_dt or pd.isna(utc_dt): return utc_dt if not pod or pod not in POD_TIMEZONE_MAPPINGS: # Keep as UTC if POD is missing or not in mappings if hasattr(utc_dt, 'tzinfo') and utc_dt.tzinfo is not None: return utc_dt.replace(tzinfo=None) return utc_dt try: # Ensure UTC timezone if hasattr(utc_dt, 'tzinfo'): if utc_dt.tzinfo is None: utc_dt = pytz.utc.localize(utc_dt) else: utc_dt = utc_dt.astimezone(pytz.utc) else: utc_dt = pytz.utc.localize(utc_dt) # Convert to POD timezone timezone_name = POD_TIMEZONE_MAPPINGS[pod] pod_tz = pytz.timezone(timezone_name) local_dt = utc_dt.astimezone(pod_tz) # Remove timezone info return local_dt.replace(tzinfo=None) except: # Fallback to UTC if conversion fails if hasattr(utc_dt, 'tzinfo') and utc_dt.tzinfo is not None: return utc_dt.replace(tzinfo=None) return utc_dt # Define all Excel columns based on the template EXCEL_COLUMNS = [ 'MAWB', 'Entry', 'POD', 'handover_doc', 'receipt_scan', 'CBP Line', 'CBP_check', 'CBP_hold', 'CBP_reject', 'CBP_abnormal', 'CBP_release', 'FDA Line', 'FDA_check', 'FDA_hold', 'FDA_release', 'FDA_reject', 'FDA_abnormal', 'CPSC Line', 'CPSC_check', 'CPSC_hold', 'CPSC_release', 'CPSC_reject', 'CPSC_abnormal', 'EPA Line', 'EPA_check', 'EPA_hold', 'EPA_release', 'EPA_reject', 'EPA_abnormal', 'USDA Line', 'USDA_check', 'USDA_hold', 'USDA_release', 'USDA_reject', 'USDA_abnormal', 'DOT Line', 'DOT_check', 'DOT_hold', 'DOT_release', 'DOT_reject', 'DOT_abnormal', 'FWS Line', 'FWS_check', 'FWS_hold', 'FWS_release', 'FWS_reject', 'FWS_abnormal', 'APHIS Line', 'APHIS_check', 'APHIS_hold', 'APHIS_release', 'APHIS_reject', 'APHIS_abnormal', 'TTB Line', 'TTB_check', 'TTB_hold', 'TTB_release', 'TTB_reject', 'TTB_abnormal', 'released', 'clearance_customs', 'truck_shipment' ] def __init__(self, notion_integration: NotionIntegration): """Initialize with Notion integration instance.""" self.notion = notion_integration def parse_mawb_input(self, mawb_input: str) -> List[str]: """Parse MAWB input string into a list of MAWBs.""" if not mawb_input or not mawb_input.strip(): return [] # Split by common separators and clean mawb_list = [] for separator in ['\n', '\r\n', ',', ';', '\t']: if separator in mawb_input: parts = mawb_input.split(separator) mawb_list.extend([part.strip() for part in parts if part.strip()]) break else: # Single MAWB mawb_list = [mawb_input.strip()] # Remove duplicates while preserving order unique_mawbs = [] seen = set() for mawb in mawb_list: if mawb and mawb not in seen: unique_mawbs.append(mawb) seen.add(mawb) return unique_mawbs def extract_entry_data(self, entry_record: Dict, pod: Optional[str] = None) -> Dict: """Extract relevant data from a Notion entry record with timezone conversion.""" properties = entry_record.get("properties", {}) data = {} # Extract MAWB if "mawb" in properties and properties["mawb"].get("title"): data["mawb"] = properties["mawb"]["title"][0].get("text", {}).get("content", "") # Extract Entry Number if "entry_number" in properties and properties["entry_number"].get("rich_text"): data["entry_number"] = properties["entry_number"]["rich_text"][0].get("text", {}).get("content", "") # Extract dates with timezone conversion date_fields = [ "ams_submit_date", "entry_submit_date", "cbp_hold_date", "cbp_release_date", "cpsc_review_date", "cpsc_release_date" ] for field in date_fields: if field in properties and properties[field].get("date"): date_str = properties[field]["date"].get("start") if date_str: try: dt = datetime.fromisoformat(date_str.replace('Z', '+00:00')) dt_naive = dt.replace(tzinfo=None) # Convert to POD timezone data[field] = self.convert_utc_to_pod_timezone(dt_naive, pod) except: data[field] = date_str else: data[field] = None else: data[field] = None # Extract text fields text_fields = ["cbp_hold_line", "cpsc_review_line"] for field in text_fields: if field in properties and properties[field].get("rich_text"): data[field] = properties[field]["rich_text"][0].get("text", {}).get("content", "") else: data[field] = "" return data def calculate_final_release(self, cbp_release_date, cpsc_review_date, cpsc_release_date): """Calculate final release date based on CBP and CPSC dates.""" if not cbp_release_date: return None if not cpsc_review_date: return cbp_release_date if not cpsc_release_date: return None if isinstance(cbp_release_date, datetime) and isinstance(cpsc_release_date, datetime): return max(cbp_release_date, cpsc_release_date) elif isinstance(cbp_release_date, datetime): return cbp_release_date elif isinstance(cpsc_release_date, datetime): return cpsc_release_date else: try: cbp_dt = datetime.fromisoformat(str(cbp_release_date).replace('Z', '+00:00')) cpsc_dt = datetime.fromisoformat(str(cpsc_release_date).replace('Z', '+00:00')) result = max(cbp_dt, cpsc_dt) return result.replace(tzinfo=None) except: return cbp_release_date def create_excel_row(self, entry_data: Dict, mawb_info: Dict) -> Dict: """Create a row for Excel export from entry and MAWB data.""" row = {col: None for col in self.EXCEL_COLUMNS} row['MAWB'] = entry_data.get("mawb", "") row['Entry'] = entry_data.get("entry_number", "") row['POD'] = mawb_info.get("pod", "") if entry_data.get("entry_submit_date"): row['handover_doc'] = entry_data["entry_submit_date"] if entry_data.get("cbp_release_date"): row['CBP_release'] = entry_data["cbp_release_date"] if entry_data.get("cbp_hold_line"): row['CBP Line'] = entry_data["cbp_hold_line"] if entry_data.get("cpsc_review_line"): row['CPSC Line'] = entry_data["cpsc_review_line"] if entry_data.get("cpsc_review_date"): row['CPSC_check'] = entry_data["cpsc_review_date"] if entry_data.get("cpsc_release_date"): row['CPSC_release'] = entry_data["cpsc_release_date"] # Map FDA data if entry_data.get("fda_review_line"): row['FDA Line'] = entry_data["fda_review_line"] if entry_data.get("fda_review_date"): row['FDA_check'] = entry_data["fda_review_date"] if entry_data.get("fda_hold_date"): row['FDA_hold'] = entry_data["fda_hold_date"] if entry_data.get("fda_release_date"): row['FDA_release'] = entry_data["fda_release_date"] final_release = self.calculate_final_release( entry_data.get("cbp_release_date"), entry_data.get("cpsc_review_date"), entry_data.get("cpsc_release_date") ) if final_release: row['released'] = final_release return row def export_to_excel(self, mawb_input: str = None, output_file: str = None, date_filter: Dict = None) -> str: """ Export entry milestone data to Excel file. Args: mawb_input: MAWB input string (optional) output_file: Output file path date_filter: Created date filter with 'start_date' and/or 'end_date' keys Returns: Path to the generated Excel file """ # Get entries if mawb_input: mawb_list = self.parse_mawb_input(mawb_input) if not mawb_list: raise ValueError("No valid MAWBs provided") entries = self.notion.get_entries_by_mawb_list(mawb_list) else: # Get all entries with date filter query_data = {} if date_filter: filter_conditions = [] if date_filter.get('start_date'): filter_conditions.append({ "timestamp": "created_time", "created_time": {"on_or_after": date_filter['start_date']} }) if date_filter.get('end_date'): filter_conditions.append({ "timestamp": "created_time", "created_time": {"on_or_before": date_filter['end_date']} }) if len(filter_conditions) == 2: query_data["filter"] = {"and": filter_conditions} elif len(filter_conditions) == 1: query_data["filter"] = filter_conditions[0] # Query with pagination all_results = [] has_more = True start_cursor = None while has_more: if start_cursor: query_data["start_cursor"] = start_cursor response = self.notion._make_database_query("entry_record", query_data) if response: all_results.extend(response.get("results", [])) has_more = response.get("has_more", False) start_cursor = response.get("next_cursor") else: break entries = all_results if not entries: raise ValueError("No entry records found for the provided criteria") # Extract MAWB list mawb_list = [] for entry in entries: entry_data = self.extract_entry_data(entry) mawb = entry_data.get("mawb", "") if mawb and mawb not in mawb_list: mawb_list.append(mawb) # Get MAWB info mawb_info = self.notion.get_mawb_info_by_list(mawb_list) if mawb_list else {} # Process data excel_rows = [] for entry in entries: # Get MAWB to look up POD mawb_for_lookup = "" if "mawb" in entry.get("properties", {}) and entry["properties"]["mawb"].get("title"): mawb_for_lookup = entry["properties"]["mawb"]["title"][0].get("text", {}).get("content", "") # Get POD from mawb_info pod = mawb_info.get(mawb_for_lookup, {}).get("pod", None) if mawb_for_lookup else None # Extract entry data with POD timezone conversion entry_data = self.extract_entry_data(entry, pod) mawb = entry_data.get("mawb", "") if mawb in mawb_info: excel_row = self.create_excel_row(entry_data, mawb_info[mawb]) excel_rows.append(excel_row) else: excel_row = self.create_excel_row(entry_data, {}) excel_rows.append(excel_row) # Create DataFrame df = pd.DataFrame(excel_rows, columns=self.EXCEL_COLUMNS) # Export to Excel with pd.ExcelWriter(output_file, engine='openpyxl') as writer: df.to_excel(writer, index=False, sheet_name='Entry Milestone Export') # Get worksheet for formatting worksheet = writer.sheets['Entry Milestone Export'] workbook = writer.book # Import styles from openpyxl.styles import NamedStyle, Font, PatternFill # Header formatting header_format = NamedStyle(name="header_format") header_format.font = Font(bold=True, color="FFFFFF") header_format.fill = PatternFill(start_color="4472C4", end_color="4472C4", fill_type="solid") # Apply header formatting for col_num, value in enumerate(df.columns.values, 1): cell = worksheet.cell(1, col_num) cell.style = header_format # Auto-adjust column widths for idx, col in enumerate(df.columns, 1): max_length = max(len(str(col)), 12) for val in df[col].astype(str): if len(val) > max_length: max_length = len(val) worksheet.column_dimensions[worksheet.cell(1, idx).column_letter].width = min(max_length + 2, 50) return output_file