import os
import re
import sqlite3
import datetime

# Database path
DB_PATH = '/workspace/scratch/travelocity_lexicon.db'
SCHEMA_PATH = '/workspace/artifacts/sec_forensic_schema.sql'
LOG_PATH = '/workspace/scratch/chang_liu_redacted_logs.txt'

# 1. Define Sample Raw Logs to populate if file does not exist
SAMPLE_LOG_DATA = """[2026-02-15 09:12:04] PLATFORM: LINE | SENDER: Chang Liu | RECIPIENT: Alyssa Peng (Apple CE Engineering)
MSG: Hey Alyssa, hope you are doing well. We are building out the new hardware team at OpenAI with Tang Tan.

[2026-02-15 09:15:22] PLATFORM: LINE | SENDER: Alyssa Peng | RECIPIENT: Chang Liu (OpenAI Technical Staff)
MSG: Hey Chang! Yes, I heard Tang Tan joined. How is it over there?

[2026-02-18 14:22:10] PLATFORM: LINE | SENDER: Chang Liu | RECIPIENT: Alyssa Peng (Apple CE Engineering)
MSG: It's amazing. We have access to incredible compute resources. I can always give you some fun side projects if you want to collaborate.

[2026-03-02 10:05:41] PLATFORM: LINE | SENDER: Chang Liu | RECIPIENT: Alyssa Peng (Apple CE Engineering)
MSG: LOL, I found out I can access the [network storage], so funny. My Apple iCloud is still syncing and the network credentials weren't revoked on the old Mac.

[2026-03-02 10:11:15] PLATFORM: LINE | SENDER: Alyssa Peng | RECIPIENT: Chang Liu (OpenAI Technical Staff)
MSG: Wait, really? Is that secure?

[2026-03-04 11:32:00] PLATFORM: LINE | SENDER: Chang Liu | RECIPIENT: Alyssa Peng (Apple CE Engineering)
MSG: Just make sure you copy the display shields and battery housing CAD files onto an external physical drive first. That way we avoid trouble with the security team and there are no network triggers. Let's keep talking on LINE here to avoid detection. Self-delete is active.

[2026-03-15 16:45:12] PLATFORM: iMessage | SENDER: Former Coworker (Apple) | RECIPIENT: Chang Liu (OpenAI Technical Staff)
MSG: Hi, this is highly irregular, please remove me from this thread. You can keep my icloud connected if you still need more time but don't pull me into this.
"""

def setup_database():
    """Initializes the database using the SEC schema if it exists."""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # Check if table 'employees' exists
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='employees';")
    if not cursor.fetchone():
        print("[*] Initializing database with SEC Forensic Schema...")
        if os.path.exists(SCHEMA_PATH):
            with open(SCHEMA_PATH, 'r') as f:
                schema_sql = f.read()
            cursor.executescript(schema_sql)
            conn.commit()
            print("[+] Schema applied successfully.")
        else:
            print("[-] Schema SQL not found, creating baseline tables manually...")
            # Fallback table creation
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS employees (
                    employee_id VARCHAR(50) PRIMARY KEY,
                    first_name VARCHAR(100) NOT NULL,
                    last_name VARCHAR(100) NOT NULL,
                    former_apple_title VARCHAR(150),
                    former_apple_division VARCHAR(100),
                    termination_date DATE,
                    openai_onboarding_date DATE,
                    current_openai_title VARCHAR(150),
                    current_openai_division VARCHAR(100),
                    line_username VARCHAR(100),
                    is_named_defendant BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS encrypted_channels (
                    channel_id VARCHAR(50) PRIMARY KEY,
                    sender_employee_id VARCHAR(50) REFERENCES employees(employee_id),
                    recipient_name VARCHAR(150) NOT NULL,
                    recipient_status VARCHAR(100) NOT NULL,
                    platform_name VARCHAR(50) NOT NULL,
                    group_chat_name VARCHAR(200),
                    message_count INTEGER DEFAULT 0,
                    self_delete_timer_seconds INTEGER DEFAULT 0,
                    earliest_message_timestamp TIMESTAMP,
                    latest_message_timestamp TIMESTAMP,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            conn.commit()
    
    # 2. Seed baseline employees to prevent Foreign Key constraint issues
    cursor.execute("INSERT OR IGNORE INTO employees (employee_id, first_name, last_name, former_apple_title, former_apple_division, termination_date, openai_onboarding_date, current_openai_title, current_openai_division, line_username, is_named_defendant) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                   ('EMP_LIU_01', 'Chang', 'Liu', 'Senior Systems Electrical Engineer', 'CE Design Group', '2026-01-22', '2026-01-23', 'Technical Staff', 'Consumer Devices Unit', 'cliu_secure', True))
    cursor.execute("INSERT OR IGNORE INTO employees (employee_id, first_name, last_name, former_apple_title, former_apple_division, termination_date, openai_onboarding_date, current_openai_title, current_openai_division, line_username, is_named_defendant) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                   ('EMP_PENG_01', 'Yu-Ting', 'Peng', 'Hardware/SoC Engineer', 'Core Engineering', '2026-04-10', '2026-04-11', 'Technical Staff', 'io Products division', 'alyssa_p', True))
    conn.commit()
    conn.close()

def parse_and_ingest():
    """Parses raw logs and updates/inserts records into the encrypted_channels table."""
    if not os.path.exists(LOG_PATH):
        print(f"[*] Creating sample text message log file at {LOG_PATH}...")
        with open(LOG_PATH, 'w') as f:
            f.write(SAMPLE_LOG_DATA)

    print(f"[*] Reading logs from {LOG_PATH}...")
    with open(LOG_PATH, 'r') as f:
        log_content = f.read()

    # Regex to parse the structured lines
    # Matches: [2026-02-15 09:12:04] PLATFORM: LINE | SENDER: Chang Liu | RECIPIENT: Alyssa Peng (Apple CE Engineering)
    # Followed by: MSG: ...
    pattern = r"\[(?P<timestamp>[^\]]+)\] PLATFORM:\s*(?P<platform>[^|]+)\|\s*SENDER:\s*(?P<sender>[^|]+)\|\s*RECIPIENT:\s*(?P<recipient>[^(\n]+)(?:\((?P<status>[^)]+)\))?\s*\nMSG:\s*(?P<msg>.*)"
    
    matches = re.finditer(pattern, log_content)
    
    # We aggregate messages by channel (Sender ID + Recipient + Platform)
    # Key: (sender_id, recipient_name, platform, status)
    # Value: { "count": X, "earliest": Y, "latest": Z }
    channels_agg = {}
    
    for match in matches:
        gd = match.groupdict()
        timestamp_str = gd['timestamp'].strip()
        platform = gd['platform'].strip()
        sender = gd['sender'].strip()
        recipient = gd['recipient'].strip()
        status = gd['status'].strip() if gd['status'] else "Unknown"
        message = gd['msg'].strip()
        
        # Map sender name to employee_id
        sender_id = None
        if "Chang Liu" in sender:
            sender_id = 'EMP_LIU_01'
        elif "Alyssa Peng" in sender or "Yu-Ting" in sender:
            sender_id = 'EMP_PENG_01'
        else:
            # Create general employee record if sender is an employee but not yet seeded
            sender_id = f"EMP_COWORKER_{abs(hash(sender)) % 1000}"
            conn = sqlite3.connect(DB_PATH)
            c = conn.cursor()
            first, *last = sender.split()
            last_name = " ".join(last) if last else "Unknown"
            c.execute("INSERT OR IGNORE INTO employees (employee_id, first_name, last_name, former_apple_title, former_apple_division, termination_date, openai_onboarding_date, current_openai_title, current_openai_division) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                      (sender_id, first, last_name, 'Engineer', 'Apple CE Group', '2026-01-01', '2026-01-02', 'Staff', 'Unknown'))
            conn.commit()
            conn.close()

        # Build unique aggregate channel key
        channel_key = (sender_id, recipient, platform, status)
        
        dt = datetime.datetime.strptime(timestamp_str, "%Y-%m-%d %H:%M:%S")
        
        if channel_key not in channels_agg:
            channels_agg[channel_key] = {
                "count": 0,
                "earliest": dt,
                "latest": dt
            }
        
        channels_agg[channel_key]["count"] += 1
        if dt < channels_agg[channel_key]["earliest"]:
            channels_agg[channel_key]["earliest"] = dt
        if dt > channels_agg[channel_key]["latest"]:
            channels_agg[channel_key]["latest"] = dt

    # 3. Insert or update in database
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    print("[*] Ingesting aggregated channel metrics into 'encrypted_channels' table...")
    
    for key, metrics in channels_agg.items():
        sender_id, recipient_name, platform, status = key
        
        # Create deterministic channel_id
        clean_platform = platform.lower().replace(" ", "")
        clean_recipient = recipient_name.lower().replace(" ", "_")
        channel_id = f"CH_{sender_id}_{clean_platform}_{clean_recipient}"
        
        # Self-delete indicators (e.g., if platform is LINE/Signal, set baseline timer)
        self_delete_timer = 3600 if platform.lower() in ['line', 'signal'] else 0
        
        # Insert or Update (Upsert in SQLite)
        cursor.execute("""
            INSERT INTO encrypted_channels (
                channel_id, sender_employee_id, recipient_name, recipient_status, 
                platform_name, message_count, self_delete_timer_seconds, 
                earliest_message_timestamp, latest_message_timestamp
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(channel_id) DO UPDATE SET
                message_count = message_count + excluded.message_count,
                earliest_message_timestamp = MIN(earliest_message_timestamp, excluded.earliest_message_timestamp),
                latest_message_timestamp = MAX(latest_message_timestamp, excluded.latest_message_timestamp);
        """, (
            channel_id, sender_id, recipient_name, status,
            platform, metrics["count"], self_delete_timer,
            metrics["earliest"].strftime("%Y-%m-%d %H:%M:%S"),
            metrics["latest"].strftime("%Y-%m-%d %H:%M:%S")
        ))
        print(f"[+] Upserted Channel: {channel_id} | Recipient: {recipient_name} | Messages: {metrics['count']}")
        
    conn.commit()
    conn.close()
    print("[+] Ingestion run completed successfully.")

if __name__ == "__main__":
    setup_database()
    parse_and_ingest()