#!/usr/bin/env python3
"""
Enhanced Inventory Import Script
Imports additional data from inventory_restructured.ods:
- MAC addresses (to communications table)
- PC-to-machine control relationships
- Asset criticality and compliance data
- VLAN, location, security scan data
"""

import csv
import sys
import pymysql
from datetime import datetime

DB_CONFIG = {
    'host': 'localhost',
    'user': '570005354',
    'password': '570005354',
    'database': 'shopdb',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}

def connect_db():
    """Connect to MySQL database"""
    try:
        conn = pymysql.connect(**DB_CONFIG)
        print(f"✓ Connected to database: {DB_CONFIG['database']}")
        return conn
    except Exception as e:
        print(f"✗ Database connection failed: {e}")
        sys.exit(1)

def get_comtype_id(cursor, typename):
    """Get communication type ID"""
    cursor.execute("SELECT comstypeid FROM comstypes WHERE typename = %s", (typename,))
    result = cursor.fetchone()
    return result['comstypeid'] if result else None

def get_relationship_type_id(cursor, relationship_type):
    """Get relationship type ID or create if doesn't exist"""
    cursor.execute("SELECT relationshiptypeid FROM relationshiptypes WHERE relationshiptype = %s", (relationship_type,))
    result = cursor.fetchone()

    if result:
        return result['relationshiptypeid']

    # Create new relationship type
    cursor.execute("INSERT INTO relationshiptypes (relationshiptype, isactive) VALUES (%s, 1)", (relationship_type,))
    return cursor.lastrowid

def import_mac_addresses(cursor, conn, rows):
    """Import MAC addresses for equipment into communications table"""
    print("\n" + "="*70)
    print("IMPORTING MAC ADDRESSES")
    print("="*70)

    network_interface_type = get_comtype_id(cursor, 'Network_Interface')
    if not network_interface_type:
        print("✗ Network_Interface communication type not found!")
        return 0

    imported = 0
    skipped = 0

    for row in rows:
        if row.get('Row_Type') != 'Equipment':
            continue

        machine_number = row.get('Machine_Number', '').strip()
        if not machine_number:
            continue

        # Get machine ID
        cursor.execute("SELECT machineid FROM machines WHERE machinenumber = %s", (machine_number,))
        machine = cursor.fetchone()
        if not machine:
            continue

        machine_id = machine['machineid']

        # Get IP and MAC addresses from inventory
        ip1 = row.get('IP Address', '').strip()
        ip2 = row.get('IP Address (interface 2)', '').strip()
        ip3 = row.get('IP Address (Interface 3)', '').strip()

        mac1 = row.get('MAC Address (interface1)', '').strip()
        mac2 = row.get('MAC address (Interface 2)', '').strip()
        mac3 = row.get('MAC Address (Interface 3)', '').strip()

        # Create communications records for each interface with data
        interfaces = [
            (ip1, mac1, 'Interface 1'),
            (ip2, mac2, 'Interface 2'),
            (ip3, mac3, 'Interface 3')
        ]

        for ip_addr, mac_addr, interface_name in interfaces:
            if ip_addr or mac_addr:  # Only create if we have at least IP or MAC
                # Check if communication already exists for this machine/IP
                cursor.execute(
                    """SELECT comid FROM communications
                       WHERE machineid = %s AND address = %s AND comstypeid = %s""",
                    (machine_id, ip_addr or '', network_interface_type)
                )
                existing = cursor.fetchone()

                if not existing:
                    # Insert new communication record
                    cursor.execute(
                        """INSERT INTO communications
                           (machineid, comstypeid, address, macaddress, interfacename, isactive)
                           VALUES (%s, %s, %s, %s, %s, 1)""",
                        (machine_id, network_interface_type, ip_addr or None, mac_addr or None,
                         interface_name)
                    )
                    conn.commit()

                    if imported < 10:  # Show first 10
                        print(f"✓ {machine_number}: {interface_name} - IP: {ip_addr or 'N/A'}, MAC: {mac_addr or 'N/A'}")
                    imported += 1
                else:
                    skipped += 1

    if imported > 10:
        print(f"... and {imported - 10} more")

    print(f"\n✓ Imported {imported} communication records, skipped {skipped} existing")
    return imported

def import_pc_machine_relationships(cursor, conn, rows):
    """Import PC-to-machine control relationships"""
    print("\n" + "="*70)
    print("IMPORTING PC-TO-MACHINE CONTROL RELATIONSHIPS")
    print("="*70)

    controls_type_id = get_relationship_type_id(cursor, 'Controls')
    print(f"✓ Relationship type 'Controls' ID: {controls_type_id}")

    imported = 0
    skipped = 0
    errors = 0

    for row in rows:
        if row.get('Row_Type') != 'Equipment':
            continue

        machine_number = row.get('Machine_Number', '').strip()
        pc_hostname = row.get('Associated_PC_Hostname', '').strip()

        if not machine_number or not pc_hostname:
            continue

        # Get equipment machine ID
        cursor.execute("SELECT machineid FROM machines WHERE machinenumber = %s", (machine_number,))
        equipment = cursor.fetchone()
        if not equipment:
            continue

        equipment_id = equipment['machineid']

        # Find PC by hostname (could be in hostname OR machinenumber for PCs)
        cursor.execute(
            """SELECT machineid FROM machines
               WHERE (hostname = %s OR machinenumber = %s)
               AND pctypeid IS NOT NULL""",
            (pc_hostname, pc_hostname)
        )
        pc = cursor.fetchone()

        if not pc:
            errors += 1
            if errors <= 5:
                print(f"  ⚠️  PC '{pc_hostname}' not found for machine {machine_number}")
            continue

        pc_id = pc['machineid']

        # Check if relationship already exists
        cursor.execute(
            """SELECT relationshipid FROM machinerelationships
               WHERE machineid = %s AND related_machineid = %s AND relationshiptypeid = %s""",
            (pc_id, equipment_id, controls_type_id)
        )
        existing = cursor.fetchone()

        if not existing:
            # Create relationship: PC controls Equipment
            cursor.execute(
                """INSERT INTO machinerelationships
                   (machineid, related_machineid, relationshiptypeid, isactive)
                   VALUES (%s, %s, %s, 1)""",
                (pc_id, equipment_id, controls_type_id)
            )
            conn.commit()

            if imported < 10:
                print(f"✓ {pc_hostname} controls {machine_number}")
            imported += 1
        else:
            skipped += 1

    if imported > 10:
        print(f"... and {imported - 10} more")

    if errors > 5:
        print(f"  ⚠️  ... and {errors - 5} more PCs not found")

    print(f"\n✓ Imported {imported} control relationships, skipped {skipped} existing, {errors} errors")
    return imported

def import_compliance_data(cursor, conn, rows):
    """Import compliance and asset criticality data"""
    print("\n" + "="*70)
    print("IMPORTING COMPLIANCE DATA")
    print("="*70)

    imported = 0
    updated = 0

    for row in rows:
        if row.get('Row_Type') != 'Equipment':
            continue

        machine_number = row.get('Machine_Number', '').strip()
        if not machine_number:
            continue

        # Get machine ID
        cursor.execute("SELECT machineid FROM machines WHERE machinenumber = %s", (machine_number,))
        machine = cursor.fetchone()
        if not machine:
            continue

        machine_id = machine['machineid']

        # Get compliance data from inventory
        third_party_managed = row.get('3rd Party (Other) Managed (Y/N)', '').strip()
        third_party_manager = row.get('3rd Party Manager', '').strip()
        ot_environment = row.get('OT Environment', '').strip()
        dod_asset_type = row.get('DoD Asset Type', '').strip()

        # Map third party managed to enum
        third_party_enum = None
        if third_party_managed.upper() == 'Y':
            third_party_enum = 'Yes'
        elif third_party_managed.upper() == 'N':
            third_party_enum = 'No'
        else:
            third_party_enum = 'NA'

        # Check if compliance record exists
        cursor.execute("SELECT complianceid FROM compliance WHERE machineid = %s", (machine_id,))
        existing = cursor.fetchone()

        if existing:
            # Update existing record
            cursor.execute(
                """UPDATE compliance SET
                   is_third_party_managed = %s,
                   third_party_manager = %s,
                   ot_asset_system = %s,
                   ot_asset_device_type = %s,
                   lastupdated = NOW()
                   WHERE complianceid = %s""",
                (third_party_enum, third_party_manager or None,
                 ot_environment or None, dod_asset_type or None,
                 existing['complianceid'])
            )
            updated += 1
        else:
            # Insert new compliance record
            cursor.execute(
                """INSERT INTO compliance
                   (machineid, is_third_party_managed, third_party_manager,
                    ot_asset_system, ot_asset_device_type)
                   VALUES (%s, %s, %s, %s, %s)""",
                (machine_id, third_party_enum, third_party_manager or None,
                 ot_environment or None, dod_asset_type or None)
            )
            imported += 1

        conn.commit()

    print(f"✓ Imported {imported} new compliance records, updated {updated} existing")
    return imported + updated

def import_security_scans(cursor, conn, rows):
    """Import security scan data"""
    print("\n" + "="*70)
    print("IMPORTING SECURITY SCAN DATA")
    print("="*70)

    imported = 0

    for row in rows:
        if row.get('Row_Type') != 'Equipment':
            continue

        machine_number = row.get('Machine_Number', '').strip()
        scan_date_str = row.get('Security_Scan_Date', '').strip()
        scan_result = row.get('Security_Scan_Result', '').strip()

        if not machine_number or not scan_date_str:
            continue

        # Get machine ID
        cursor.execute("SELECT machineid FROM machines WHERE machinenumber = %s", (machine_number,))
        machine = cursor.fetchone()
        if not machine:
            continue

        machine_id = machine['machineid']

        # Parse scan date
        try:
            scan_date = datetime.strptime(scan_date_str, '%Y-%m-%d')
        except:
            try:
                scan_date = datetime.strptime(scan_date_str, '%m/%d/%Y')
            except:
                continue

        # Map scan result to enum
        result_map = {
            'PASS': 'Pass',
            'FAIL': 'Fail',
            'WARNING': 'Warning',
            'INFO': 'Info'
        }
        mapped_result = result_map.get(scan_result.upper(), 'Info') if scan_result else 'Info'

        # Check if this scan already exists
        cursor.execute(
            """SELECT scanid FROM compliancescans
               WHERE machineid = %s AND scan_date = %s""",
            (machine_id, scan_date)
        )
        existing = cursor.fetchone()

        if not existing:
            # Insert new scan record
            cursor.execute(
                """INSERT INTO compliancescans
                   (machineid, scan_date, scan_result, scan_details, scan_name)
                   VALUES (%s, %s, %s, %s, %s)""",
                (machine_id, scan_date, mapped_result,
                 f"Imported from inventory on {datetime.now().strftime('%Y-%m-%d')}",
                 'Inventory Import Scan')
            )
            conn.commit()
            imported += 1

    print(f"✓ Imported {imported} security scan records")
    return imported

def main():
    """Main import process"""
    print("="*70)
    print("ENHANCED INVENTORY DATA IMPORT")
    print("Importing: MAC addresses, PC-to-machine relationships,")
    print("           compliance data, and security scans")
    print("="*70)
    print()

    # Connect to database
    conn = connect_db()
    cursor = conn.cursor()

    # Read CSV file
    csv_file = '/tmp/inventory_restructured.csv'
    print(f"📄 Reading inventory file: {csv_file}")

    try:
        with open(csv_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            rows = list(reader)
        print(f"✓ Loaded {len(rows)} rows from inventory")
        print()
    except Exception as e:
        print(f"✗ Failed to read CSV file: {e}")
        conn.close()
        sys.exit(1)

    # Import each data type
    mac_count = import_mac_addresses(cursor, conn, rows)
    relationship_count = import_pc_machine_relationships(cursor, conn, rows)
    compliance_count = import_compliance_data(cursor, conn, rows)
    scan_count = import_security_scans(cursor, conn, rows)

    # Summary
    print("\n" + "="*70)
    print("IMPORT SUMMARY")
    print("="*70)
    print(f"MAC addresses imported:           {mac_count}")
    print(f"PC-to-machine relationships:      {relationship_count}")
    print(f"Compliance records:               {compliance_count}")
    print(f"Security scan records:            {scan_count}")
    print(f"Total records created:            {mac_count + relationship_count + compliance_count + scan_count}")
    print()

    # Final verification queries
    print("="*70)
    print("VERIFICATION")
    print("="*70)

    cursor.execute(
        """SELECT COUNT(*) as count FROM communications c
           JOIN machines m ON c.machineid = m.machineid
           WHERE m.pctypeid IS NULL"""
    )
    equipment_comms = cursor.fetchone()['count']
    print(f"Equipment with communications:    {equipment_comms}")

    cursor.execute(
        """SELECT COUNT(*) as count FROM machinerelationships mr
           JOIN relationshiptypes rt ON mr.relationshiptypeid = rt.relationshiptypeid
           WHERE rt.relationshiptype = 'Controls'"""
    )
    control_relationships = cursor.fetchone()['count']
    print(f"PC control relationships:         {control_relationships}")

    cursor.execute("SELECT COUNT(*) as count FROM compliance WHERE is_third_party_managed IS NOT NULL")
    third_party_count = cursor.fetchone()['count']
    print(f"Machines with compliance data:    {third_party_count}")

    cursor.execute("SELECT COUNT(*) as count FROM compliancescans")
    total_scans = cursor.fetchone()['count']
    print(f"Total security scans:             {total_scans}")

    print()

    # Close connection
    cursor.close()
    conn.close()
    print("✓ Enhanced import complete!")

if __name__ == '__main__':
    main()
