#!/usr/bin/env python3
"""
Iran OSINT Database Manager
Usage: python db_manager.py [command] [args]

Commands:
    list [table]     - List all records from a table
    add_person       - Add a person interactively
    add_hash         - Add email/hash pair
    search [term]    - Search across all tables
    export           - Export all data to JSON
"""

import sqlite3
import json
import sys
import os

DB_PATH = os.path.join(os.path.dirname(__file__), '..', 'iran_osint.db')

def get_connection():
    return sqlite3.connect(DB_PATH)

def list_table(table_name):
    """List all records from a table"""
    conn = get_connection()
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    try:
        c.execute(f"SELECT * FROM {table_name}")
        rows = c.fetchall()

        print(f"\n{'='*60}")
        print(f"    TABLE: {table_name} ({len(rows)} records)")
        print("="*60)

        for row in rows:
            print(dict(row))
            print("-"*40)
    except Exception as e:
        print(f"Error: {e}")
    finally:
        conn.close()

def add_person(name, title=None, org=None, role=None, status='Active', notes=None):
    """Add a person to the database"""
    conn = get_connection()
    c = conn.cursor()

    c.execute('''INSERT INTO persons (name, title, organization, role, status, notes)
                 VALUES (?, ?, ?, ?, ?, ?)''',
              (name, title, org, role, status, notes))

    conn.commit()
    print(f"[+] Added person: {name}")
    conn.close()

def add_hash(email, hash_val, source=None):
    """Add email hash to database (create table if needed)"""
    conn = get_connection()
    c = conn.cursor()

    # Create hashes table if not exists
    c.execute('''CREATE TABLE IF NOT EXISTS email_hashes (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        email TEXT,
        md5_hash TEXT,
        source TEXT,
        gravatar_exists INTEGER DEFAULT 0,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    )''')

    c.execute('''INSERT INTO email_hashes (email, md5_hash, source)
                 VALUES (?, ?, ?)''', (email, hash_val, source))

    conn.commit()
    print(f"[+] Added hash: {email} -> {hash_val}")
    conn.close()

def search_all(term):
    """Search across all tables"""
    conn = get_connection()
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    tables = ['persons', 'organizations', 'events', 'sources', 'social_media']

    print(f"\n[*] Searching for: {term}\n")

    for table in tables:
        try:
            c.execute(f"SELECT * FROM {table} WHERE * LIKE '%{term}%'")
        except:
            # Fallback: search specific columns
            c.execute(f"PRAGMA table_info({table})")
            columns = [col[1] for col in c.fetchall() if col[2] == 'TEXT']

            if columns:
                where_clause = " OR ".join([f"{col} LIKE ?" for col in columns])
                params = [f"%{term}%" for _ in columns]
                c.execute(f"SELECT * FROM {table} WHERE {where_clause}", params)

                rows = c.fetchall()
                if rows:
                    print(f"\n[+] Found in {table}:")
                    for row in rows:
                        print(f"    {dict(row)}")

    conn.close()

def export_all():
    """Export entire database to JSON"""
    conn = get_connection()
    conn.row_factory = sqlite3.Row
    c = conn.cursor()

    # Get all tables
    c.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = [row[0] for row in c.fetchall()]

    data = {}
    for table in tables:
        c.execute(f"SELECT * FROM {table}")
        rows = c.fetchall()
        data[table] = [dict(row) for row in rows]

    output_path = os.path.join(os.path.dirname(__file__), '..', 'resources', 'db_export.json')
    with open(output_path, 'w') as f:
        json.dump(data, f, indent=2, default=str)

    print(f"[+] Exported to: {output_path}")
    conn.close()

def show_stats():
    """Show database statistics"""
    conn = get_connection()
    c = conn.cursor()

    print("\n" + "="*60)
    print("    IRAN OSINT DATABASE STATISTICS")
    print("="*60)

    c.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = [row[0] for row in c.fetchall()]

    for table in tables:
        c.execute(f"SELECT COUNT(*) FROM {table}")
        count = c.fetchone()[0]
        print(f"    {table:25} : {count} records")

    conn.close()

if __name__ == "__main__":
    if len(sys.argv) < 2:
        show_stats()
        print("\nUsage: python db_manager.py [list|search|export] [args]")
    elif sys.argv[1] == "list":
        table = sys.argv[2] if len(sys.argv) > 2 else "persons"
        list_table(table)
    elif sys.argv[1] == "search":
        term = sys.argv[2] if len(sys.argv) > 2 else ""
        search_all(term)
    elif sys.argv[1] == "export":
        export_all()
    elif sys.argv[1] == "stats":
        show_stats()