#!/usr/bin/env python3

import os
import re
import sys
from dotenv import load_dotenv
from neo4j import GraphDatabase
from typing import Dict, Set

def extract_emails_from_cypher(content: str) -> Set[str]:
    """Extract all email addresses from Cypher queries that match identities."""
    emails = set()

    # Pattern to match email addresses in Identity matches
    # Only matches: { email: 'email@domain.com' } format
    email_pattern = r"email:\s*'([^']+)'"

    matches = re.findall(email_pattern, content)
    emails.update(matches)

    return emails

def get_uuid_for_email(session, email: str) -> str:
    """Query the database to get UUID for a given email address."""
    query = "MATCH (i:Identity { email: $email }) RETURN i.id"
    result = session.run(query, {"email": email})
    record = result.single()

    if record and record["i.id"]:
        return record["i.id"]
    else:
        print(f"WARNING: No UUID found for email: {email}")
        return None

def replace_email_with_uuid(content: str, email_to_uuid: Dict[str, str]) -> str:
    """Replace email-based identity matches with UUID-based ones."""
    result = content

    for email, uuid in email_to_uuid.items():
        if uuid is None:
            continue

        # First, add comment after --changeset line for this specific email
        # Find changesets that contain this specific email
        changeset_pattern = re.compile(rf"(--changeset [^\n]*\n)(?=(?:(?!--changeset).)*email:\s*'{re.escape(email)}')", re.DOTALL)
        changeset_replacement = rf"\1// email: {email}\n"
        result = changeset_pattern.sub(changeset_replacement, result)

        # Replace patterns like: MATCH (i:Identity { email: 'email@domain.com' })
        pattern1 = re.compile(rf"(\(\s*i\s*:\s*Identity\s*\{{\s*)email:\s*'{re.escape(email)}'(\s*\}})")
        replacement1 = rf"\1id: '{uuid}'\2"
        result = pattern1.sub(replacement1, result)

        # Replace patterns in rollback sections too
        pattern2 = re.compile(rf"(--rollback.*?\(\s*i\s*:\s*Identity\s*\{{\s*)email:\s*'{re.escape(email)}'(\s*\}})", re.DOTALL)
        replacement2 = rf"\1id: '{uuid}'\2"
        result = pattern2.sub(replacement2, result)

    return result

def main():
    # Load environment variables from .env file
    load_dotenv()

    cypher_file = "PLATFORM-4586-vendor-star-cleanup.cypher"

    try:
        # Read the original Cypher file
        with open(cypher_file, 'r', encoding='utf-8') as f:
            content = f.read()

        # Extract all email addresses
        emails = extract_emails_from_cypher(content)

        if not emails:
            print("No email addresses found in the Cypher file.")
            return

        print(f"Found {len(emails)} unique email addresses:")
        for email in sorted(emails):
            print(f"  - {email}")

        # Connect to Neo4j database
        uri = os.getenv('NEO4J_URL')
        user = os.getenv('NEO4J_USERNAME')
        password = os.getenv('NEO4J_PASSWORD')

        if not all([uri, user, password]):
            print("Missing required Neo4j environment variables")
            print("Make sure .env file contains NEO4J_URL, NEO4J_USERNAME, and NEO4J_PASSWORD")
            return

        driver = GraphDatabase.driver(uri, auth=(user, password))

        try:
            driver.verify_connectivity()
            print("Successfully connected to Neo4j database")
        except Exception as e:
            print(f"Failed to connect to Neo4j database: {e}")
            return

        # Query database for UUIDs
        email_to_uuid = {}

        print("\nQuerying database for UUIDs...")
        with driver.session() as session:
            for email in sorted(emails):
                try:
                    uuid = get_uuid_for_email(session, email)
                    email_to_uuid[email] = uuid
                    if uuid:
                        print(f"  {email} -> {uuid}")
                    else:
                        print(f"  {email} -> NOT FOUND")
                except Exception as e:
                    print(f"Error querying email {email}: {e}")
                    email_to_uuid[email] = None

        driver.close()

        # Check if we found all UUIDs
        missing_uuids = [email for email, uuid in email_to_uuid.items() if uuid is None]
        if missing_uuids:
            print(f"\nWARNING: Could not find UUIDs for {len(missing_uuids)} emails:")
            for email in missing_uuids:
                print(f"  - {email}")
            print("\nThe script will continue but these emails will not be replaced.")

        # Replace emails with UUIDs
        print("\nReplacing emails with UUIDs in Cypher file...")
        updated_content = replace_email_with_uuid(content, email_to_uuid)

        # Write the updated content to a new file
        output_file = cypher_file.replace('.cypher', '-with-uuids.cypher')
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(updated_content)

        print(f"Updated Cypher file written to: {output_file}")

        # Show summary
        successful_replacements = len([uuid for uuid in email_to_uuid.values() if uuid is not None])
        print(f"\nSummary:")
        print(f"  - Total emails found: {len(emails)}")
        print(f"  - Successful UUID lookups: {successful_replacements}")
        print(f"  - Failed UUID lookups: {len(missing_uuids)}")

    except FileNotFoundError:
        print(f"Error: Could not find file '{cypher_file}'")
        print("Make sure the file exists in the current directory.")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()