#!/usr/bin/env python3
import sqlite3
import geoip2.database

def get_location(ip_address, reader):
    """
    Given an IP address, return the location in 'City, Country' format.
    If the lookup fails, returns 'Unknown, Unknown'.
    """
    try:
        response = reader.city(ip_address)
        city = response.city.name if response.city.name else "Unknown"
        country = response.country.name if response.country.name else "Unknown"
        return city, country
    except Exception:
        return "Unknown, Unknown"

def main():
    # Initialize the GeoLite2 reader
    reader = geoip2.database.Reader('GeoLite2-City.mmdb')
    
    # Connect to your SQLite database (update the path if necessary)
    conn = sqlite3.connect('access_log.db')
    conn.row_factory = sqlite3.Row  # Enable name-based access to columns
    cursor = conn.cursor()
    
    # Create a new table with the updated structure (location instead of ip_address)
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS file_access_log_new (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT,
            rel_path TEXT,
            filesize INTEGER,
            mime TEXT,
            city TEXT,
            country TEXT,
            user_agent TEXT,
            device_id TEXT,
            cached BOOLEAN
        )
    ''')
    conn.commit()
    
    # Read all rows from the old table
    cursor.execute('SELECT * FROM file_access_log')
    rows = cursor.fetchall()
    
    # Insert rows into the new table, converting IP addresses to locations
    for row in rows:
        timestamp = row['timestamp']
        rel_path = row['rel_path']
        filesize = row['filesize']
        mime = row['mime']
        ip_address = row['ip_address']
        city, country = get_location(ip_address, reader)
        user_agent = row['user_agent']
        device_id = row['device_id']
        cached = row['cached']
        
        cursor.execute('''
            INSERT INTO file_access_log_new
            (timestamp, rel_path, filesize, mime, city, country, user_agent, device_id, cached)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (timestamp, rel_path, filesize, mime, city, country, user_agent, device_id, cached))
    
    conn.commit()
    
    # Optional: Replace the old table with the new table
    cursor.execute('DROP TABLE file_access_log')
    cursor.execute('ALTER TABLE file_access_log_new RENAME TO file_access_log')
    conn.commit()
    
    # Clean up: close connections
    conn.close()
    reader.close()
    
    print("Database conversion complete.")

if __name__ == '__main__':
    main()