import pymysql
import json
import csv
import sys

def main():
    try:
        with open('config.json', 'r') as f:
            config = json.load(f)
    except FileNotFoundError:
        print("Error: config.json not found.")
        sys.exit(1)
        
    databases = config.get('databases', [])
    if not databases:
        print("No databases to check.")
        return
        
    host = config['host']
    port = config.get('port', 3306)
    user = config['user']
    password = config['password']
    table = config['table']
    image_column = config.get('image_column', 'test_image_path')
    
    summary_data = []
    
    print(f"Starting to count rows and columns in {len(databases)} databases...")
    
    for db_name in databases:
        try:
            connection = pymysql.connect(
                host=host,
                user=user,
                password=password,
                database=db_name,
                port=port,
                connect_timeout=5
            )
            
            with connection.cursor() as cursor:
                # Count rows where image is present
                query_rows = f"SELECT COUNT(*) as count FROM {table} WHERE {image_column} IS NOT NULL AND {image_column} != ''"
                cursor.execute(query_rows)
                row_count = cursor.fetchone()[0]
                
                # Count columns in the table
                query_cols = f"SELECT count(*) FROM information_schema.columns WHERE table_schema = '{db_name}' AND table_name = '{table}'"
                cursor.execute(query_cols)
                col_count = cursor.fetchone()[0]
                
                print(f"[{db_name}] => {row_count} rows, {col_count} columns")
                summary_data.append([db_name, row_count, col_count, "Success"])
                
            connection.close()
            
        except pymysql.err.OperationalError as e:
            # Usually means database doesn't exist or connection failed
            print(f"[{db_name}] => Database not found or connection failed: {e.args[1]}")
            summary_data.append([db_name, 0, 0, f"Error: DB not found ({e.args[1]})"])
            
        except pymysql.err.ProgrammingError as e:
            # Usually means table or column doesn't exist
            print(f"[{db_name}] => Table or column missing: {e.args[1]}")
            summary_data.append([db_name, 0, 0, f"Error: Table/Column missing ({e.args[1]})"])
            
        except Exception as e:
            print(f"[{db_name}] => Error: {e}")
            summary_data.append([db_name, 0, 0, f"Error: {e}"])
            
    # Write summary report
    summary_filename = "DATABASE_ROW_AND_COLUMN_COUNTS.csv"
    with open(summary_filename, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(['Database Name', 'Row Count', 'Column Count', 'Status'])
        for row in summary_data:
            writer.writerow(row)
            
    print(f"\n=====================================")
    print(f"ALL DONE! Counts saved to: {summary_filename}")
    print(f"=====================================")

if __name__ == '__main__':
    main()
