class SQLViewer:
    def __init__(self, db_connection_method):
        self.execute = db_connection_method

    def get_schema_info(self, schema='public'):
        # Get all table names in the specified schema
        tables_result = self.execute("""
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = %s AND table_type = 'BASE TABLE';
        """, [schema])
        tables = [row['table_name'] for row in tables_result]

        schema_info = {}

        for table in tables:
            # Get columns and data types
            columns_result = self.execute("""
                SELECT column_name, data_type
                FROM information_schema.columns
                WHERE table_schema = %s AND table_name = %s;
            """, [schema, table])
            columns = [(row['column_name'], row['data_type']) for row in columns_result]

            # Get foreign keys
            foreign_keys_result = self.execute("""
                SELECT
                    kcu.column_name AS fk_column,
                    ccu.table_name AS referenced_table,
                    ccu.column_name AS referenced_column
                FROM
                    information_schema.table_constraints AS tc
                    JOIN information_schema.key_column_usage AS kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                    JOIN information_schema.constraint_column_usage AS ccu
                    ON ccu.constraint_name = tc.constraint_name
                    AND ccu.table_schema = tc.table_schema
                WHERE
                    tc.constraint_type = 'FOREIGN KEY' AND
                    tc.table_schema = %s AND
                    tc.table_name = %s;
            """, [schema, table])
            foreign_keys = [
                (row['fk_column'], row['referenced_table'], row['referenced_column'])
                for row in foreign_keys_result
            ]

            schema_info[table] = {
                'columns': columns,
                'foreign_keys': foreign_keys
            }

        return schema_info
    
    def map_data_type(self, postgres_type):
        type_mapping = {
            'integer': 'int',
            'bigint': 'int',
            'smallint': 'int',
            'character varying': 'string',
            'varchar': 'string',
            'text': 'string',
            'date': 'date',
            'timestamp without time zone': 'datetime',
            'timestamp with time zone': 'datetime',
            'boolean': 'bool',
            'numeric': 'float',
            'real': 'float'
            # Add more mappings as needed
        }
        return type_mapping.get(postgres_type, 'string')  # Default to 'string' if type not mapped

    def generate_mermaid_er(self, schema_info):
        mermaid_lines = ["erDiagram"]
        
        for table, info in schema_info.items():
            # Define the table and its columns
            mermaid_lines.append(f"    {table} {{")
            for column_name, data_type in info['columns']:
                # Convert PostgreSQL data types to Mermaid-compatible types
                mermaid_data_type = self.map_data_type(data_type)
                mermaid_lines.append(f"        {mermaid_data_type} {column_name}")
            mermaid_lines.append("    }")
        
        # Define relationships
        for table, info in schema_info.items():
            for fk_column, referenced_table, referenced_column in info['foreign_keys']:
                # Mermaid relationship syntax: [Table1] }|--|| [Table2] : "FK_name"
                relation = f"    {table} }}|--|| {referenced_table} : \"{fk_column} to {referenced_column}\""
                mermaid_lines.append(relation)
        
        return "\n".join(mermaid_lines)