#!/usr/bin/env python3
"""
Schema extractor - comprehensive database schema analysis.
Enhanced from tool_runner_3a with precision_orchestrator's detailed requirements.
"""

import sqlite3
import json
from pathlib import Path

def extract_schema(conn):
    """Extract complete schema information."""
    cursor = conn.cursor()
    
    # Get database overview
    cursor.execute("SELECT COUNT(*) FROM sqlite_master WHERE type='table'")
    table_count = cursor.fetchone()[0]
    
    # Get all tables with their SQL
    cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table' ORDER BY name")
    tables_info = cursor.fetchall()
    
    tables = {}
    total_rows = 0
    
    for table_name, create_sql in tables_info:
        # Get column information
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns_raw = cursor.fetchall()
        
        columns = []
        primary_keys = []
        for col in columns_raw:
            column_info = {
                "name": col[1],
                "type": col[2],
                "nullable": not col[3],
                "default": col[4],
                "primary_key": bool(col[5])
            }
            columns.append(column_info)
            if col[5]:
                primary_keys.append(col[1])
        
        # Get foreign keys
        cursor.execute(f"PRAGMA foreign_key_list({table_name})")
        foreign_keys_raw = cursor.fetchall()
        
        foreign_keys = []
        for fk in foreign_keys_raw:
            foreign_keys.append({
                "column": fk[3],
                "references_table": fk[2],
                "references_column": fk[4]
            })
        
        # Get indexes
        cursor.execute(f"PRAGMA index_list({table_name})")
        indexes_raw = cursor.fetchall()
        
        indexes = []
        for idx in indexes_raw:
            cursor.execute(f"PRAGMA index_info({idx[1]})")
            idx_cols = cursor.fetchall()
            indexes.append({
                "name": idx[1],
                "unique": bool(idx[2]),
                "columns": [col[2] for col in idx_cols]
            })
        
        # Get row count
        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
        row_count = cursor.fetchone()[0]
        total_rows += row_count
        
        # Sample data for small tables
        sample_data = []
        if row_count > 0 and row_count <= 1000:
            cursor.execute(f"SELECT * FROM {table_name} LIMIT 5")
            sample_rows = cursor.fetchall()
            col_names = [col[0] for col in cursor.description]
            sample_data = [dict(zip(col_names, row)) for row in sample_rows]
        
        tables[table_name] = {
            "row_count": row_count,
            "columns": columns,
            "primary_keys": primary_keys,
            "foreign_keys": foreign_keys,
            "indexes": indexes,
            "create_sql": create_sql,
            "sample_data": sample_data
        }
    
    # Determine database complexity
    if table_count <= 5:
        complexity = "Simple"
    elif table_count <= 15:
        complexity = "Medium"
    else:
        complexity = "Complex"
    
    # Identify main entities
    main_entities = []
    for table_name, info in tables.items():
        if info["row_count"] > 100 and not any(fk for fk in info["foreign_keys"]):
            main_entities.append(table_name)
    
    overview = {
        "table_count": table_count,
        "total_rows": total_rows,
        "complexity": complexity,
        "main_entities": main_entities[:5]  # Top 5 main entities
    }
    
    return {
        "overview": overview,
        "tables": tables
    }

def main():
    """Main entry point."""
    db_path = Path("./database.sqlite")
    
    if not db_path.exists():
        print(json.dumps({"error": "Database file not found"}))
        return
    
    try:
        conn = sqlite3.connect(str(db_path))
        results = extract_schema(conn)
        conn.close()
        
        # Output as JSON
        print(json.dumps(results, indent=2))
        
    except Exception as e:
        print(json.dumps({"error": str(e)}))

if __name__ == "__main__":
    main()