# Add parent directory to Python path so we can import helpers
import sys
from pathlib import Path

parent_dir = Path(__file__).parent.parent
sys.path.insert(0, str(parent_dir))

import psycopg2
import csv
from datetime import datetime
from decimal import Decimal
from proteus import Model

from helpers.config import (
    PURCHASE_CONTRACTS_CSV, 
    connect_to_tryton,
    get_db_connection,
    DB_CONFIG  # Add this to your config
)

from helpers.tryton_helpers import (
    
    parse_decimal, 
    parse_date,
    ensure_party_is_supplier,
    find_party_by_name,
    find_uom_by_code,
    find_currency_by_code,
    find_warehouse,
    find_location,
    find_payment_term_by_name,
    find_product_by_code,
    find_incoterm_by_code,
    find_weight_basis_by_name,
    get_party_invoice_address,
    find_purchase_contract_by_number,
    find_or_create_analytic_dimension_value,
    link_analytic_dimensions_to_purchase,
    ensure_party_has_category
)

# Import migration mapping helper
from helpers.migration_mapping import MigrationMapper


# CSV Configuration
CSV_FILE_PATH = PURCHASE_CONTRACTS_CSV 


# Default values
DEFAULT_STATE = 'draft'
DEFAULT_INVOICE_METHOD = 'manual'
DEFAULT_INVOICE_STATE = 'none'
DEFAULT_SHIPMENT_STATE = 'none'

# Import options
AUTO_ENABLE_SUPPLIERS = True  # Set to False to skip auto-enabling supplier flag
SKIP_NON_SUPPLIERS = False   # Set to True to skip parties that aren't suppliers


def update_purchase_custom_fields(purchase_id, custom_data):
    """Update custom fields in purchase using direct SQL"""
    if not custom_data:
        return True
    
    conn = get_db_connection()
    if not conn:
        print(f"  ⚠ Could not update custom fields - database connection failed")
        return False
    
    try:
        cursor = conn.cursor()
        
        # First, check what columns exist in purchase_purchase table
        cursor.execute("""
            SELECT column_name, data_type 
            FROM information_schema.columns 
            WHERE table_name = 'purchase_purchase'
            AND column_name IN ('number')
            ORDER BY column_name
        """)
        existing_columns = {row[0]: row[1] for row in cursor.fetchall()}
        
        print(f"    Available custom columns in purchase_purchase:")
        for col_name, col_type in existing_columns.items():
            print(f"      - {col_name} ({col_type})")
        
        # Build UPDATE query for custom fields
        set_clauses = []
        values = []
        
        if 'number' in existing_columns and custom_data.get('number'):
            set_clauses.append("number = %s")
            values.append(custom_data['number'])
            print(f"    Adding number = {custom_data['number']}")
        
        if set_clauses:
            values.append(purchase_id)
            update_query = f"""
                UPDATE purchase_purchase 
                SET {', '.join(set_clauses)}
                WHERE id = %s
            """
            
            print(f"    Executing UPDATE with fields: {', '.join([c.split('=')[0].strip() for c in set_clauses])}")
            
            cursor.execute(update_query, values)
            rows_affected = cursor.rowcount
            conn.commit()
            
            if rows_affected > 0:
                print(f"    ✓ {rows_affected} row(s) updated successfully")
            else:
                print(f"    ⚠ No rows updated (purchase_id={purchase_id} not found?)")
        else:
            print(f"    No custom fields to update (either no data provided or columns not found)")
        
        cursor.close()
        conn.close()
        return True
        
    except Exception as e:
        print(f"  ⚠ Error updating custom fields: {e}")
        import traceback
        traceback.print_exc()
        if conn:
            conn.rollback()
            conn.close()
        return False


def update_line_custom_fields(line_id, custom_data):
    """Update custom fields in purchase line using direct SQL"""
    if not custom_data:
        return True
    
    conn = get_db_connection()
    if not conn:
        print(f"  ⚠ Could not update line custom fields - database connection failed")
        return False
    
    try:
        cursor = conn.cursor()
        
        # Build UPDATE query for custom fields
        set_clauses = []
        values = []
        
        if custom_data.get('from_del'):
            set_clauses.append("from_del = %s")
            values.append(custom_data['from_del'])
        
        if custom_data.get('to_del'):
            set_clauses.append("to_del = %s")
            values.append(custom_data['to_del'])
        
        if set_clauses:
            values.append(line_id)
            update_query = f"""
                UPDATE purchase_line 
                SET {', '.join(set_clauses)}
                WHERE id = %s
            """
            
            cursor.execute(update_query, values)
            conn.commit()
        
        cursor.close()
        conn.close()
        return True
        
    except Exception as e:
        print(f"  ⚠ Error updating line custom fields: {e}")
        import traceback
        traceback.print_exc()
        if conn:
            conn.rollback()
            conn.close()
        return False


def create_pricing_estimated(line_id, pricing_data):
    """Create pricing_estimated record using direct SQL"""
    if not pricing_data or not pricing_data.get('trigger'):
        return True
    
    conn = get_db_connection()
    if not conn:
        print(f"  ⚠ Could not create pricing estimate - database connection failed")
        return False
    
    try:
        cursor = conn.cursor()
        
        # Check if pricing_estimated table exists and what its structure is
        cursor.execute("""
            SELECT column_name 
            FROM information_schema.columns 
            WHERE table_name = 'pricing_estimated'
            ORDER BY ordinal_position
        """)
        
        columns = [row[0] for row in cursor.fetchall()]
        
        if not columns:
            print(f"    Info: pricing_estimated table does not exist, skipping pricing estimate")
            cursor.close()
            conn.close()
            return True
        
        # Insert into pricing_estimated table
        insert_query = """
            INSERT INTO pricing_estimated (line, trigger, estimated_date, create_date, write_date, create_uid, write_uid)
            VALUES (%s, %s, %s, NOW(), NOW(), 1, 1)
        """
        
        cursor.execute(insert_query, (
            line_id,
            pricing_data['trigger'],
            pricing_data.get('estimated_date')
        ))
        
        conn.commit()
        cursor.close()
        conn.close()
        print(f"    ✓ Pricing estimate created successfully")
        return True
        
    except psycopg2.errors.ForeignKeyViolation as e:
        # This is expected if pricing_estimated references purchase_line only
        print(f"    Info: Pricing estimate skipped (table references purchase_line only, not sale_line)")
        if conn:
            conn.rollback()
            conn.close()
        return True  # Return True to continue processing
        
    except Exception as e:
        print(f"  ⚠ Error creating pricing estimate: {e}")
        import traceback
        traceback.print_exc()
        if conn:
            conn.rollback()
            conn.close()
        return False


def import_purchases(csv_file):
    """Import purchases from CSV file with migration mapping tracking"""
    Purchase = Model.get('purchase.purchase')
    PurchaseLine = Model.get('purchase.line')
    
    print(f"{'='*70}")
    print(f"IMPORTING PURCHASES FROM CSV")
    print(f"{'='*70}\n")
    print(f"Reading from: {csv_file}\n")
    
    imported_count = 0
    skipped_count = 0
    error_count = 0
    errors = []

    # Get company (assuming single company or default)
    Company = Model.get('company.company')
    companies = Company.find([])
    if not companies:
        print("✗ Error: No company found in the system")
        return
    company = companies[0]
    print(f"Using company: {company.rec_name}\n") 
    
    # Collect all mappings for batch insert at the end
    purchase_mappings = []
    line_mappings = []
    
    try:
        # Initialize migration mapper
        with MigrationMapper(DB_CONFIG) as mapper:
            
            with open(csv_file, 'r', encoding='utf-8-sig') as file:
                reader = csv.DictReader(file)
                row_num = 0
                
                for row in reader:
                    row_num += 1
                    
                    try:
                        # Extract fields from CSV
                        source_purchase_id = row.get('source_id', '').strip()  # Source system ID
                        number = row.get('number', '').strip()
                        reference = row.get('reference', '').strip()
                        our_reference = row.get('our_reference', '').strip()
                        
                        
                        if not number:
                            print(f"Row {row_num}: Skipping - no number\n")
                            continue
                        
                        print(f"{'='*70}")
                        print(f"Row {row_num}: Processing purchase {number}")
                        print(f"{'='*70}")
                        
                        # CHECK IF ALREADY IMPORTED using migration mapper
                        if source_purchase_id:
                            existing_tryton_id = mapper.get_tryton_id('purchase_contract', source_purchase_id)
                            if existing_tryton_id:
                                print(f"  ⏭ Purchase already imported (Source ID: {source_purchase_id} -> Tryton ID: {existing_tryton_id})")
                                skipped_count += 1
                                print()
                                continue
                        
                        # Alternative: Check by number using existing helper
                        existing_purchase = find_purchase_contract_by_number(number)
                        if existing_purchase:
                            print(f"  ⏭ Purchase {number} already exists (ID: {existing_purchase.id})")
                            # Save mapping even if it already exists (for reconciliation)
                            if source_purchase_id:
                                purchase_mappings.append({
                                    'object_type': 'purchase_contract',
                                    'source_id': source_purchase_id,
                                    'tryton_model': 'purchase.purchase',
                                    'tryton_id': existing_purchase.id,
                                    'recon_key': number
                                })
                            skipped_count += 1
                            print()
                            continue
                        
                        # Parse other fields
                        raw_purchase_date = row.get('purchase_date', '').strip()
                        try:
                            purchase_date = datetime.strptime(raw_purchase_date, '%m/%d/%Y').date()
                        except (ValueError, AttributeError):
                            purchase_date = parse_date(raw_purchase_date)
                        party_name = row.get('party_name', '').strip()
                        
                        # Find related records
                        party = find_party_by_name(party_name)
                        if not party:
                            raise ValueError(f"Party not found: {party_name}")
                        
                        # Check party is supplier
                        if not ensure_party_is_supplier(party, auto_enable=AUTO_ENABLE_SUPPLIERS):
                            if SKIP_NON_SUPPLIERS:
                                print(f"  ⏭ Skipping - party {party_name} is not a supplier\n")
                                skipped_count += 1
                                continue
                            else:
                                raise ValueError(f"Party {party_name} is not a supplier")
                        
                        # Reload party after category addition to get fresh data
                        Party = Model.get('party.party')
                        party = Party(party.id)

                        # Find invoice address
                        invoice_address = get_party_invoice_address(party)
                        if not invoice_address:
                            raise ValueError(f"No invoice address found for party {party_name}")
                        
                        # Parse additional fields
                        currency = find_currency_by_code(row.get('currency', 'USD'))
                        warehouse = find_warehouse(row.get('warehouse'))
                        payment_term = find_payment_term_by_name(row.get('payment_term'))
                        weight_basis_abbr = find_weight_basis_by_name(row.get('weight_basis'))
                        tol_min_pct = parse_decimal(row.get('tol_min_pct', ''), 'tol_min_pct')
                        tol_max_pct = parse_decimal(row.get('tol_max_pct', ''), 'tol_max_pct')
                        tol_min_qty = parse_decimal(row.get('tol_min_qty', ''), 'tol_min_qty')
                        tol_max_qty = parse_decimal(row.get('tol_max_qty', ''), 'tol_max_qty')                        
                        
                        from_location_name = row.get('from_location_name', '').strip()
                        from_location = find_location(from_location_name)
                        
                        to_location_name = row.get('to_location_name', '').strip()
                        to_location = find_location(to_location_name)   

                        incoterm_code = row.get('incoterm_name', '').strip()
                        incoterm = find_incoterm_by_code(incoterm_code, 2025)    

                        description = row.get('description', '').strip()
                        comment = row.get('comment', '').strip()        
                        trader_name = row.get('trader', '').strip()
                        operator_name = row.get('operator', '').strip()
                        

                        
                        demurrage = parse_decimal(row.get('demurrage', ''), 'demurrage')
                        laytime_hours = parse_decimal(row.get('laytime_hours', ''), 'laytime_hours')
                        nor_extra_hours = parse_decimal(row.get('nor_extra_hours', ''), 'nor_extra_hours')
                        pumping_rate = parse_decimal(row.get('pumping_rate', ''), 'pumping_rate')

                                                
                        # CREATE PURCHASE
                        print(f"  Creating purchase...")
                        purchase= Purchase()
                        purchase.company = company
                        purchase.reference = reference
                        purchase.our_reference = our_reference
                        purchase.party = party
                        purchase.invoice_address = invoice_address
                        
                        purchase.purchase_date = purchase_date
                        purchase.currency = currency
                        if warehouse:
                            purchase.warehouse = warehouse
                        purchase.payment_term = payment_term
                        purchase.wb = weight_basis_abbr
                        purchase.tol_min = tol_min_pct
                        purchase.tol_max = tol_max_pct
                        purchase.tol_min_qt = tol_min_qty
                        purchase.tol_max_qt = tol_max_qty
                        purchase.incoterm = incoterm
                        purchase.from_location = from_location
                        purchase.to_location = to_location
                        purchase.description = description
                        purchase.comment = comment
                        purchase.state = DEFAULT_STATE
                        purchase.invoice_method = DEFAULT_INVOICE_METHOD

                        # Retrieve trader (optional)
                        if trader_name:
                            trader = find_party_by_name(trader_name)
                            if not trader:
                                raise ValueError(f"Trader not found: {trader_name}")
                            trader, _ = ensure_party_has_category(trader, 'TRADER')
                            purchase.trader = Party(trader.id)

                        # Retrieve operator (optional)
                        if operator_name:
                            operator = find_party_by_name(operator_name)
                            if not operator:
                                raise ValueError(f"Operator not found: {operator_name}")
                            operator, _ = ensure_party_has_category(operator, 'OPERATOR')
                            purchase.operator = Party(operator.id)


                        # Save the purchase
                        purchase.save()
                        print(f"  ✓ Purchase created (ID: {purchase.id})")
                        
                        # Update custom fields (like 'number')
                        custom_fields = {'number': number}
                        update_purchase_custom_fields(purchase.id, custom_fields)
                        
                        # SAVE MIGRATION MAPPING for purchase
                        if source_purchase_id:
                            purchase_mappings.append({
                                'object_type': 'purchase_contract',
                                'source_id': source_purchase_id,
                                'tryton_model': 'purchase.purchase',
                                'tryton_id': purchase.id,
                                'recon_key': number
                            })
                            print(f"  📝 Mapping queued: Source {source_purchase_id} -> Tryton {purchase.id}")
                        
                        # Process purchase lines
                        line_type = row.get('line_type', 'line').strip()
                        source_line_id = row.get('source_line_id', '').strip()
                        
                        if line_type == 'line':
                            # Regular product line
                            product_code = row.get('line_product_code', '').strip()
                            quantity = parse_decimal(row.get('line_quantity', ''), 'quantity')
                            unit_price = parse_decimal(row.get('line_price', ''), 'unit_price')
                            period_at = row.get('period_at', '').strip().lower()
                            concentration = parse_decimal(row.get('concentration', ''), 'concentration')

                            product = find_product_by_code(product_code)
                            if not product:
                                raise ValueError(f"Product not found: {product_code}")
                            
                            unit = find_uom_by_code(row.get('line_unit_code', ''))

                            # Parse shipping dates
                            raw_from_del = row.get('line_from_del', '').strip()
                            try:
                                from_del = datetime.strptime(raw_from_del, '%m/%d/%Y').date()
                            except (ValueError, AttributeError):
                                from_del = parse_date(raw_from_del)

                            raw_to_del = row.get('line_to_del', '').strip()
                            try:
                                to_del = datetime.strptime(raw_to_del, '%m/%d/%Y').date()
                            except (ValueError, AttributeError):
                                to_del = parse_date(raw_to_del)                            
                            
                            # Create line
                            line = PurchaseLine()
                            line.purchase = purchase
                            line.type = 'line'
                            sequence = 1 # Default sequence, can be enhanced to handle multiple lines
                            line.sequence = sequence  
                            line.product = product
                            line.quantity = quantity
                            line.unit = unit if unit else product.purchase_uom
                            line.unit_price = unit_price
                            line.from_del = from_del
                            line.to_del = to_del
                            line.period_at = period_at
                            line.concentration = concentration

                            
                            # Optional fields
                            description = row.get('description', '').strip()
                            if description:
                                line.description = description

                            line.save()
                            
                            # # Update line custom fields
                            # line_custom = {}
                            # from_del = row.get('from_del', '').strip()
                            # to_del = row.get('to_del', '').strip()
                            # if from_del:
                            #     line_custom['from_del'] = from_del
                            # if to_del:
                            #     line_custom['to_del'] = to_del
                            
                            # if line_custom:
                            #     update_line_custom_fields(line.id, line_custom)
                            
                            # Create pricing estimate if applicable
                            pricing_trigger = row.get('pricing_trigger', '').strip()
                            raw_pricing_date = row.get('pricing_estimated_date', '').strip()
                            try:
                                pricing_estimated_date = datetime.strptime(raw_pricing_date, '%m/%d/%Y').date()
                            except (ValueError, AttributeError):
                                pricing_estimated_date = parse_date(raw_pricing_date)
                            if pricing_trigger:
                                pricing_data = {
                                    'trigger': pricing_trigger,
                                    'estimated_date': pricing_estimated_date
                                }
                                create_pricing_estimated(line.id, pricing_data)
                            
                            # SAVE MIGRATION MAPPING for line
                            if source_line_id:
                                line_mappings.append({
                                    'object_type': 'purchase_line',
                                    'source_id': source_line_id,
                                    'tryton_model': 'purchase.line',
                                    'tryton_id': line.id,
                                    'recon_key': f"{number}-Line {sequence}-{product_code}"
                                })
                            
                            print(f"  ✓ Added line (ID: {line.id})")
                            print(f"    Product: {product.rec_name}")
                            print(f"    Quantity: {quantity}")
                            
                        else:
                            # Comment, subtitle, or other line types
                            line = PurchaseLine()
                            line.purchase = purchase
                            line.type = line_type
                            line.description = row.get('description', '').strip()
                            line.save()
                            
                            print(f"  ✓ Added {line_type} line (ID: {line.id})")
                        
                        
                        # **NEW: Handle analytic dimensions**
                        book_name = row.get('book', '').strip()
                        strategy_name = row.get('strategy', '').strip()

                        print(f"\nProcessing analytic dimensions...")

                        # Collect all dimension values in one dictionary
                        dimension_values = {}
                    
                        # Handle Book dimension
                        if book_name:
                            book_dimension_value = find_or_create_analytic_dimension_value('Book', book_name)
                            if book_dimension_value:
                                dimension_values['Book'] = book_dimension_value
                            else:
                                print(f"  ⚠ Could not find/create Book dimension value: {book_name}")

                        # Handle Strategy dimension
                        if strategy_name:
                            strategy_dimension_value = find_or_create_analytic_dimension_value('Strategy', strategy_name)
                            if strategy_dimension_value:
                                dimension_values['Strategy'] = strategy_dimension_value
                            else:
                                print(f"  ⚠ Could not find/create Strategy dimension value: {strategy_name}")
                                                

                        # Link ALL dimensions to purchase in ONE call
                        if dimension_values:
                            link_analytic_dimensions_to_purchase(purchase, dimension_values)
                        else:
                            print(f"  No analytic dimensions to link")


                        imported_count += 1
                        print(f"✓ Successfully imported purchase {number}\n")
                        
                    except Exception as e:
                        error_msg = f"Row {row_num} - {number if 'number' in locals() else 'Unknown'}: {str(e)}"
                        errors.append(error_msg)
                        error_count += 1
                        print(f"✗ Error on row {row_num}: {e}\n")
                        import traceback
                        traceback.print_exc()
            
            # BATCH SAVE ALL MAPPINGS at the end
            print(f"\n{'='*70}")
            print("SAVING MIGRATION MAPPINGS")
            print(f"{'='*70}\n")
            
            if purchase_mappings:
                print(f"Saving {len(purchase_mappings)} purchase mappings...")
                mapper.save_mappings_batch(purchase_mappings)
                print(f"✓ Purchase mappings saved\n")
            
            if line_mappings:
                print(f"Saving {len(line_mappings)} line mappings...")
                mapper.save_mappings_batch(line_mappings)
                print(f"✓ Line mappings saved\n")
        
        # Summary
        print(f"{'='*70}")
        print("IMPORT SUMMARY")
        print(f"{'='*70}")
        print(f"Successfully imported: {imported_count} purchases")
        print(f"Skipped (already exist): {skipped_count} purchases")
        print(f"Errors: {error_count}")
        print(f"Migration mappings saved: {len(purchase_mappings)} purchases, {len(line_mappings)} lines")
        
        if errors:
            print(f"\nError details:")
            for error in errors:
                print(f"  - {error}")
        
        print(f"\n{'='*70}")
        
    except FileNotFoundError:
        print(f"✗ Error: CSV file not found at {csv_file}")
        print(f"Please update CSV_FILE_PATH in the script with the correct path.")
    except Exception as e:
        print(f"✗ Fatal error: {e}")
        import traceback
        traceback.print_exc()


def verify_import():
    """Verify imported purchases and their migration mappings"""
    Purchase = Model.get('purchase.purchase')
    
    print(f"\n{'='*70}")
    print("VERIFICATION - Purchases and Migration Mappings")
    print(f"{'='*70}\n")
    
    # Find all purchases (or limit to recently created ones)
    purchases = Purchase.find([], order=[('id', 'DESC')])
    
    if purchases:
        print(f"Found {len(purchases)} purchases (showing last 10):\n")
        print(f"{'ID':<8} {'Number':<15} {'Reference':<15} {'Party':<25} {'State':<12} {'Source ID':<15}")
        print("-" * 100)
        
        # Initialize mapper to look up source IDs
        with MigrationMapper(DB_CONFIG) as mapper:
            for purchase in purchases[:10]:  # Show last 10 created
                purchase_id = purchase.id
                
                # Get number from database since it's a custom field
                conn = get_db_connection()
                number = 'N/A'
                if conn:
                    try:
                        cursor = conn.cursor()
                        cursor.execute("SELECT number FROM purchase_purchase WHERE id = %s", (purchase_id,))
                        result = cursor.fetchone()
                        if result and result[0]:
                            number = str(result[0])[:14]
                        cursor.close()
                        conn.close()
                    except:
                        if conn:
                            conn.close()
                
                # Look up source ID from migration mapping
                source_id = 'N/A'
                try:
                    cursor = mapper.connection.cursor()
                    cursor.execute("""
                        SELECT source_id[1] 
                        FROM public.os_migration_mapping
                        WHERE tryton_id = %s 
                          AND 'purchase_contract' = ANY(object_type)
                        ORDER BY write_date DESC
                        LIMIT 1
                    """, (purchase_id,))
                    result = cursor.fetchone()
                    if result and result[0]:
                        source_id = str(result[0])[:14]
                    cursor.close()
                except:
                    pass
                
                reference = purchase.reference[:14] if purchase.reference else 'N/A'
                party = purchase.party.rec_name[:24] if purchase.party else 'N/A'
                state = purchase.state if purchase.state else 'N/A'
                
                print(f"{purchase_id:<8} {number:<15} {reference:<15} {party:<25} {state:<12} {source_id:<15}")
                
                # Show lines
                if purchase.lines:
                    print(f"  Lines: {len(purchase.lines)}")
                    for line in purchase.lines[:3]:  # Show first 3 lines
                        if line.type == 'line' and line.product:
                            print(f"    - {line.product.rec_name[:40]} | Qty: {line.quantity} | Price: {line.unit_price}")
                        else:
                            print(f"    - [{line.type}] {(line.description or '')[:40]}")
    else:
        print("No purchases found")
    
    print()


def check_mapping_stats():
    """Display statistics about migration mappings"""
    print(f"\n{'='*70}")
    print("MIGRATION MAPPING STATISTICS")
    print(f"{'='*70}\n")
    
    try:
        with MigrationMapper(DB_CONFIG) as mapper:
            cursor = mapper.connection.cursor()
            
            # Count mappings by object type
            cursor.execute("""
                SELECT 
                    object_type[1] as obj_type,
                    COUNT(*) as count,
                    MIN(write_date) as first_import,
                    MAX(write_date) as last_import
                FROM public.os_migration_mapping
                GROUP BY object_type[1]
                ORDER BY count DESC
            """)
            
            results = cursor.fetchall()
            
            if results:
                print(f"{'Object Type':<25} {'Count':<10} {'First Import':<15} {'Last Import':<15}")
                print("-" * 70)
                for row in results:
                    obj_type = row[0] or 'N/A'
                    count = row[1]
                    first = row[2].strftime('%Y-%m-%d') if row[2] else 'N/A'
                    last = row[3].strftime('%Y-%m-%d') if row[3] else 'N/A'
                    print(f"{obj_type:<25} {count:<10} {first:<15} {last:<15}")
                
                # Total count
                cursor.execute("SELECT COUNT(*) FROM public.os_migration_mapping")
                total = cursor.fetchone()[0]
                print(f"\nTotal mappings: {total}")
            else:
                print("No migration mappings found")
            
            cursor.close()
            
    except Exception as e:
        print(f"Error retrieving mapping statistics: {e}")
        import traceback
        traceback.print_exc()
    
    print()


def prepare_parties_as_supplier(csv_file):
    """Pre-process: Add SUPPLIER category to all parties in CSV using Proteus"""
    Party = Model.get('party.party')
    Category = Model.get('party.category')
    
    print(f"{'='*70}")
    print(f"PREPARING PARTIES AS SUPPLIERS (via Categories)")
    print(f"{'='*70}\n")
    
    # Find SUPPLIER category
    supplier_categories = Category.find([('name', '=', 'SUPPLIER')])
    if not supplier_categories:
        # Try case-insensitive
        all_categories = Category.find([])
        for cat in all_categories:
            if cat.name.upper() == 'SUPPLIER':
                supplier_categories = [cat]
                break
    
    if not supplier_categories:
        print(f"✗ SUPPLIER category not found in the system!")
        print(f"Please create a party category named 'SUPPLIER' first.\n")
        return False
    
    supplier_category = supplier_categories[0]
    print(f"Found SUPPLIER category (ID: {supplier_category.id})\n")
    
    party_names = set()
    
    try:
        with open(csv_file, 'r', encoding='utf-8-sig') as file:
            reader = csv.DictReader(file)
            for row in reader:
                party_name = row.get('party_name', '').strip()
                if party_name:
                    party_names.add(party_name)
        
        print(f"Found {len(party_names)} unique parties in CSV\n")
        
        updated_count = 0
        already_supplier_count = 0
        not_found_count = 0
        
        for party_name in party_names:
            print(f"Processing party: {party_name}")
            
            # Find party
            parties = Party.find([('name', '=', party_name)])
            if not parties:
                parties = Party.find([('code', '=', party_name)])
            
            if not parties:
                print(f"  ✗ Not found\n")
                not_found_count += 1
                continue
            
            party = parties[0]
            
            # Check if already has SUPPLIER category
            has_supplier = False
            if party.categories:
                for cat in party.categories:
                    if cat.name.upper() == 'SUPPLIER':
                        has_supplier = True
                        break
            
            if has_supplier:
                print(f"  ✓ Already has SUPPLIER category\n")
                already_supplier_count += 1
                continue
            
            # Add SUPPLIER category using Proteus
            try:
                # Reload party and category in same context
                party_to_update = Party(party.id)
                supplier_cat = Category(supplier_category.id)
                
                party_to_update.categories.append(supplier_cat)
                party_to_update.save()
                print(f"  ✓ SUPPLIER category added\n")
                updated_count += 1
            except Exception as e:
                print(f"  ✗ Failed: {e}\n")
        
        print(f"{'='*70}")
        print(f"PREPARATION SUMMARY")
        print(f"{'='*70}")
        print(f"Already have SUPPLIER category: {already_supplier_count}")
        print(f"SUPPLIER category added: {updated_count}")
        print(f"Not found: {not_found_count}")
        print(f"{'='*70}\n")
        
        return True
        
    except FileNotFoundError:
        print(f"✗ Error: CSV file not found at {csv_file}\n")
        return False
    except Exception as e:
        print(f"✗ Error: {e}\n")
        import traceback
        traceback.print_exc()
        return False


def main():
    print("="*70)
    print("TRYTON PURCHASE IMPORT SCRIPT WITH MIGRATION MAPPING")
    print("Using Proteus with XML-RPC Connection")
    print("="*70)
    print()
    
    # Connect to Tryton using XML-RPC
    if not connect_to_tryton():
        return 1
    
    # Optional: Prepare all parties as suppliers first
    # Uncomment the following line to mark all parties in CSV as suppliers before importing
    prepare_parties_as_supplier(CSV_FILE_PATH)
    
    # Import purchases with migration mapping
    import_purchases(CSV_FILE_PATH)
    
    # Verify import
    verify_import()
    
    # Show mapping statistics
    check_mapping_stats()
    
    return 0


if __name__ == '__main__':
    exit(main())
