entry_milestone_exporter.py

#!/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