import_to_db.py

import pandas as pd
import psycopg2
from psycopg2.extras import execute_values
import os
from dotenv import load_dotenv
import uuid
from datetime import datetime
from decimal import Decimal

# Load environment variables
load_dotenv()

def connect_to_db():
    """Create database connection"""
    return psycopg2.connect(
        os.getenv('DATABASE_URL')
    )

def process_excel_file(file_path):
    """Process Excel file with enhanced data extraction"""
    try:
        df = pd.read_excel(file_path)
        data = df.to_dict('records')
        
        processed_data = []
        for row in data:
            # Derive contact_name from first and last name
            first_name = str(row.get('First Name', '')).strip()
            last_name = str(row.get('Last Name', '')).strip()
            contact_name = f"{first_name} {last_name}".strip()
            
            # Derive company_domain from email or website
            email = str(row.get('Email', '')).strip().lower()
            website = str(row.get('Website', '')).strip().lower()
            company_domain = ''
            if '@' in email:
                company_domain = email.split('@')[1]
            elif website:
                company_domain = website.replace('http://', '').replace('https://', '').split('/')[0]
            
            contact_data = {
                # Basic Contact Info
                'id': str(uuid.uuid4()),
                'company_name': str(row.get('Company', '')).strip(),
                'company_domain': company_domain,
                'job_title': str(row.get('Title', '')).strip(),
                'contact_name': contact_name,
                'contact_first_name': first_name,
                'contact_middle_name': None,
                'contact_last_name': last_name,
                'contact_email': email,
                'contact_email_status': str(row.get('Email Status', 'unverified')).strip(),
                
                # Address Information
                'company_address_1': str(row.get('Company Address', '')).strip(),
                'company_address_2': None,
                'city': str(row.get('Company City', '')).strip() or str(row.get('City', '')).strip(),
                'state': str(row.get('Company State', '')).strip() or str(row.get('State', '')).strip(),
                'country': str(row.get('Company Country', '')).strip() or str(row.get('Country', '')).strip(),
                
                # Contact Details
                'phone': str(row.get('Company Phone', '')).strip(),
                'primary_email': str(row.get('Primary Email', '')).strip() or None,
                'catch_all_status': bool(row.get('Catch-all Status', False)),
                'corporate_phone': str(row.get('Corporate Phone', '')).strip() or None,
                'direct_dial': None,
                
                # Professional Info
                'employee_size': int(row.get('# Employees', 0)) if pd.notna(row.get('# Employees')) else 0,
                'revenue': str(row.get('Annual Revenue', '')).strip(),
                'sic': '',  # Not provided in Excel
                'industry_type': str(row.get('Industry', '')).strip(),
                'verification_results': '',
                
                # Social & Professional URLs
                'linkedin_url': str(row.get('Company Linkedin Url', '')).strip(),
                'person_linkedin_url': str(row.get('Person Linkedin Url', '')).strip() or None,
                'facebook_url': str(row.get('Facebook Url', '')).strip() or None,
                'twitter_url': str(row.get('Twitter Url', '')).strip() or None,
                'github_url': None,
                
                # Professional Context
                'seniority': str(row.get('Seniority', 'OTHER')).strip(),
                'departments': str(row.get('Departments', '')).strip().split(',') if row.get('Departments') else [],
                'skills': str(row.get('Keywords', '')).strip().split(',') if row.get('Keywords') else [],
                'languages': [],
                'account_owner': str(row.get('Account Owner', '')).strip() or None,
                'lead_score': None,
                'notes': None,
                
                # Company Profile Data
                'company_profile': {
                    'domain': company_domain,
                    'company_aboutus': str(row.get('SEO Description', '')).strip(),
                    'company_size': str(row.get('# Employees', '')).strip(),
                    'industry': str(row.get('Industry', '')).strip(),
                    'company_Link': str(row.get('Website', '')).strip(),
                    'companyBanner_url': '',
                    'company_title': str(row.get('Company', '')).strip(),
                    'headline': '',
                    'company_img_url': '',
                    'location': [],
                    'company_stocks': '',
                    'specialties': str(row.get('Technologies', '')).strip(),
                    
                    # Financial Information
                    'annual_revenue': Decimal(str(row.get('Annual Revenue', 0))) if pd.notna(row.get('Annual Revenue')) else None,
                    'total_funding': Decimal(str(row.get('Total Funding', 0))) if pd.notna(row.get('Total Funding')) else None,
                    'last_funding_type': str(row.get('Latest Funding', '')).strip() or None,
                    'last_funding_amount': Decimal(str(row.get('Latest Funding Amount', 0))) if pd.notna(row.get('Latest Funding Amount')) else None,
                    'last_funding_date': pd.to_datetime(row.get('Last Raised At')).to_pydatetime() if pd.notna(row.get('Last Raised At')) else None,
                }
            }
            processed_data.append(contact_data)
        
        return processed_data
    except Exception as e:
        print(f"Error processing Excel file {file_path}: {e}")
        return None
    
def insert_companies(conn, companies_data):
    """Insert or update companies in the database"""
    cursor = conn.cursor()
    
    # Deduplicate companies by domain, keeping the last occurrence
    unique_companies = {}
    for company in companies_data:
        unique_companies[company['domain']] = company
    
    companies_data = list(unique_companies.values())
    print(f"Attempting to insert/update {len(companies_data)} unique companies")
    
    insert_query = """
    INSERT INTO "companyProfilesData" (
        domain, company_aboutus, "companySize", industry, "company_Link",
        "companyBanner_url", company_title, headline, company_img_url,
        "createdAt", "updatedAt", location, "companyStocks", specialties,
        legal_name, trading_name, sub_industry, naics_code, sic_code,
        annual_revenue, revenue_range, profitability, stock_symbol,
        stock_exchange, market_cap, total_funding, last_funding_type,
        last_funding_amount, last_funding_date
    ) VALUES %s
    ON CONFLICT (domain) DO UPDATE SET
        company_aboutus = EXCLUDED.company_aboutus,
        "companySize" = EXCLUDED."companySize",
        industry = EXCLUDED.industry,
        "updatedAt" = CURRENT_TIMESTAMP
    RETURNING domain
    """
    
    try:
        company_values = []
        for company in companies_data:
            company_values.append((
                company['domain'],
                company['aboutus'],
                company['companySize'],
                company['industry'],
                company['companyLink'],
                company['companyBannerUrl'],
                company['title'],
                company['headline'],
                company['companyImgUrl'],
                company['createdAt'],
                company['updatedAt'],
                company['location'],
                company['companyStocks'],
                company['specialties'],
                company['legalName'],
                company['tradingName'],
                company['subIndustry'],
                company['naicsCode'],
                company['sicCode'],
                company['revenue'],  # This maps to annual_revenue in DB
                company['revenueRange'],
                company['profitability'],
                company['stockSymbol'],
                company['stockExchange'],
                company['marketCap'],
                company['totalFunding'],
                company['lastFundingType'],
                company['lastFundingAmount'],
                company['lastFundingDate']
            ))
        
        if company_values:
            execute_values(cursor, insert_query, company_values)
            domains = cursor.fetchall()
            conn.commit()
            print(f"Successfully committed {len(domains)} companies")
            return [domain[0] for domain in domains]
        else:
            print("No valid companies to insert")
            return []
            
    except Exception as e:
        conn.rollback()
        print(f"Error inserting companies: {str(e)}")
        raise
    finally:
        cursor.close()

def insert_contacts(conn, contacts_data):
    """Insert or update contacts in the database"""
    cursor = conn.cursor()
    print(f"Attempting to insert/update {len(contacts_data)} contacts")
    
    insert_query = """
    INSERT INTO contacts (
        id, company_name, company_domain, job_title, contact_name,
        contact_first_name, contact_middle_name, contact_last_name,
        contact_email, contact_email_status, company_address_1,
        company_address_2, city, state, zip, country, phone,
        employee_size, revenue, sic, industry_type, verification_results,
        linkedin_url, domain
    ) VALUES %s
    ON CONFLICT (id) DO UPDATE SET
        contact_email = EXCLUDED.contact_email,
        job_title = EXCLUDED.job_title,
        zip = EXCLUDED.zip
    """
    
    try:
        contact_values = []
        for contact in contacts_data:
            # Use empty string if no domain exists
            domain = contact['company_domain'] if contact['company_domain'] else ''
            
            # Ensure zip is never None
            zip_code = contact.get('zip', '')
            if zip_code is None:
                zip_code = ''
                
            contact_values.append((
                contact['id'],
                contact['company_name'],
                domain,
                contact['job_title'],
                contact['contact_name'],
                contact['contact_first_name'],
                contact['contact_middle_name'],
                contact['contact_last_name'],
                contact['contact_email'],
                contact['contact_email_status'],
                contact['company_address_1'],
                contact['company_address_2'],
                contact['city'],
                contact['state'],
                zip_code,  # Using the sanitized zip value
                contact['country'],
                contact['phone'],
                contact['employee_size'],
                contact['revenue'],
                contact['sic'],
                contact['industry_type'],
                contact['verification_results'],
                contact['linkedin_url'],
                domain
            ))
        
        if contact_values:
            print(f"Executing insert for {len(contact_values)} contacts")
            execute_values(cursor, insert_query, contact_values)
            conn.commit()
            print(f"Successfully committed {len(contact_values)} contacts")
        else:
            print("No valid contacts to insert")
            
    except Exception as e:
        conn.rollback()
        print(f"Error inserting contacts: {str(e)}")
        raise
    finally:
        cursor.close()
        
def process_folder(conn, folder_path):
    """Process all Excel files in a specific folder"""
    try:
        for filename in os.listdir(folder_path):
            if filename.endswith(('.xlsx', '.xls')):
                file_path = os.path.join(folder_path, filename)
                print(f"\nProcessing file: {file_path}")
                
                contacts_data = process_excel_file(file_path)
                if contacts_data:
                    print(f"Found {len(contacts_data)} contacts to process")
                    
                    # First, extract and insert company data
                    companies_data = []
                    for contact in contacts_data:
                        companies_data.append({
                            'domain': contact['company_domain'],
                            'aboutus': '',  # company_aboutus in DB
                            'companySize': 'SMALL',  # Enum value
                            'industry': contact['industry_type'],
                            'companyLink': '',  # company_Link in DB
                            'companyBannerUrl': '',  # companyBanner_url in DB
                            'title': contact['company_name'],  # company_title in DB
                            'headline': '',
                            'companyImgUrl': '',  # company_img_url in DB
                            'createdAt': datetime.now(),
                            'updatedAt': datetime.now(),
                            'location': [],
                            'companyStocks': '',
                            'specialties': '',
                            'legalName': None,
                            'tradingName': None,
                            'subIndustry': None,
                            'naicsCode': None,
                            'sicCode': contact['sic'],
                            'revenue': Decimal('0'),  # annual_revenue in DB
                            'revenueRange': None,
                            'profitability': None,
                            'stockSymbol': None,
                            'stockExchange': None,
                            'marketCap': None,
                            'totalFunding': None,
                            'lastFundingType': None,
                            'lastFundingAmount': None,
                            'lastFundingDate': None
                        })
                    
                    print(f"Found {len(companies_data)} companies to process")
                    
                    # Insert companies
                    if companies_data:
                        insert_companies(conn, companies_data)
                    
                    # Insert all contacts without filtering
                    insert_contacts(conn, contacts_data)
                else:
                    print(f"No data found in file: {file_path}")
                    
    except Exception as e:
        print(f"Error processing folder {folder_path}: {e}")
        raise

def main():
    """Main function to process files and import to database"""
    # Base directory containing the File1, File2, etc. folders
    base_dir = 'split_output'
    
    # Connect to database
    conn = connect_to_db()
    
    try:
        # Process each folder (File1, File2, etc.)
        for folder_name in os.listdir(base_dir):
            folder_path = os.path.join(base_dir, folder_name)
            
            # Check if it's a directory
            if os.path.isdir(folder_path):
                print(f"\nProcessing folder: {folder_name}")
                process_folder(conn, folder_path)
                
    except Exception as e:
        print(f"Error in main process: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    main()