import re
import csv
from typing import Dict
from pathlib import Path

from sqlalchemy import inspect, create_engine, text, event
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.engine import Engine


from db_query_mcp.db_adapters.base_adapter import BaseAdapter


__all__ = ['RelationalDBAdapter']


supported_dbs = [
    'sqlite',      # SQLite（内置支持，无需额外驱动）
    'postgresql',  # PostgreSQL (如 psycopg2)
    'mysql',       # MySQL (如 pymysql)
    'oracle',      # Oracle (如 cx_Oracle)
    'mssql',       # Microsoft SQL Server (如 pyodbc)
    'firebird',    # Firebird
    'sybase',      # Sybase
    'db2',         # IBM DB2
    'informix',    # IBM Informix
]


class RelationalDBAdapter(BaseAdapter):
    
    def __init__(self, db_uri: str):
        self.db_uri = db_uri
        self._test_uri(self.db_uri)
        self.engine = create_engine(db_uri)
        self._test_connection(self.engine)
        self._register_event_listeners(self.engine)

    def query(self, sql: str) -> Dict:
        self._check_sql(sql)

        with self.engine.connect() as conn:
            result = conn.execute(text(sql))
            columns = list(result.keys())  
            data = result.fetchall()

            return self._format_query_result_to_markdown(columns, data)

    def export(self, sql: str, output: str):
        sql = self._check_sql(sql)

        with self.engine.connect() as conn:
            result = conn.execute(text(sql))
            columns = list(result.keys())
            data = result.fetchall()

            output = Path(output)

            if output.is_dir():
                path = output / 'export.csv'
                if path.exists():
                    raise FileExistsError(f'File {path.resolve()} already exists.')
            elif output.exists():
                raise FileExistsError(f'File {output} already exists.')
            else:
                path = output

            self._export_to_file(path, columns, data)
            return f'Successfully exported the data to {path.resolve()}'

    def get_db_type(self) -> str:
        try:
            db_type = self.db_uri.split(':')[0].lower()
            if db_type not in supported_dbs:
                raise ValueError(f'Unsupported database type: {db_type}')
        except Exception:
            raise ValueError(f'Your db uri is not valid: {self.db_uri}')

        return db_type

    def get_db_schema(self, include_foreign_keys: bool = True) -> Dict[str, Dict]:
        schema = {}
        try:
            inspector = inspect(self.engine)
            tables = inspector.get_table_names()
        
            for table_name in tables:
                table_info = {
                    'columns': [],
                    'primary_key': [],
                    'foreign_keys': []
                }
            
                for column in inspector.get_columns(table_name):
                    table_info['columns'].append({
                        'name': column['name'],
                        'type': column['type'],
                        'nullable': column['nullable'],
                        'default': column['default'],
                        'comment': column.get('comment', '')
                    })
            
                pk_info = inspector.get_pk_constraint(table_name)
                table_info['primary_key'] = pk_info.get('constrained_columns', [])
            
                if include_foreign_keys:
                    for fk in inspector.get_foreign_keys(table_name):
                        table_info['foreign_keys'].append({
                            'constrained_columns': fk['constrained_columns'],
                            'referred_table': fk['referred_table'],
                            'referred_columns': fk['referred_columns']
                        })
            
                # Add table comment if available
                try:
                    table_comment = inspector.get_table_comment(table_name)
                    if table_comment and table_comment.get('text'):
                        table_info['comment'] = table_comment['text']
                    else:
                        table_info['comment'] = ''
                except (NotImplementedError, AttributeError):
                    # Some databases don't support table comments
                    table_info['comment'] = ''
            
                schema[table_name] = table_info
            
            schema = self._format_schema_to_markdown(schema)

            return schema
        
        except SQLAlchemyError as e:
            raise SQLAlchemyError(f'Failed to obtain the database schema: {str(e)}')

    def _test_uri(self, db_uri: str):
        db_type = db_uri.split(':')[0].lower().strip()
        if db_type not in supported_dbs:
            raise ValueError(f'Unsupported database type: {db_type}')
        
        if db_type == 'sqlite':
            path = Path(db_uri.split('///')[1].strip())
            if not path.exists():
                raise FileNotFoundError(f'Database file {path.resolve()} not found.')

    def _test_connection(self, engine: Engine):
        try:
            with engine.connect() as conn:
                conn.execute(text('SELECT 1'))
        except SQLAlchemyError as e:
            raise SQLAlchemyError(f'Failed to connect to the database: {str(e)}')

    def _register_event_listeners(self, engine: Engine):
        @event.listens_for(engine, "before_execute")
        def prevent_write_operations(conn, clauseelement, multiparams, params):
            if hasattr(clauseelement, 'is_dml') and clauseelement.is_dml:
                raise Exception("DML operations are not allowed. Only SELECT queries are permitted.")
            if hasattr(clauseelement, 'is_ddl') and clauseelement.is_ddl:
                raise Exception("DDL operations are not allowed. Only SELECT queries are permitted.")

    def _check_sql(self, sql: str) -> str:
        sql = sql.strip()

        if not sql.upper().startswith('SELECT'):
            raise Exception('Only support query operations.') 

        forbidden_pattern = r'\b(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER)\b'
        if re.search(forbidden_pattern, sql, re.IGNORECASE):
            raise Exception(f'Only support query operations.')

        return  sql

    def _export_to_file(self, output: str, columns: list, data: list):
        with open(output, 'w', encoding='utf-8', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(columns)
            writer.writerows(data)

    def _format_schema_to_markdown(self, schema: Dict[str, Dict]) -> str:
        markdown_output = []
        markdown_output.append('## Database Schema\n')
    
        for table_name, table_info in schema.items():
            # Table header
            markdown_output.append(f'### Table name: `{table_name}`\n')
            
            # Table comment if available
            if table_info.get('comment'):
                markdown_output.append(f'**Table Description:** {table_info["comment"]}\n')
        
            # Columns table
            if table_info.get('columns'):
                markdown_output.append('### Table columns\n')
                markdown_output.append('| Column Name | Data Type | Nullable | Default Value | Primary Key | Comment |')
                markdown_output.append('|---|---|---|---|---|---|')
            
                primary_keys = table_info.get('primary_key', [])
            
                for column in table_info['columns']:
                    name = column['name']
                    data_type = str(column['type'])
                    nullable = '✓' if column['nullable'] else '✗'
                    default = str(column['default']) if column['default'] is not None else '-'
                    is_pk = '✓' if name in primary_keys else '✗'
                    comment = column.get('comment', '') or '-'
                
                    markdown_output.append(f'| {name} | {data_type} | {nullable} | {default} | {is_pk} | {comment} |')
            
                markdown_output.append('')
        
            # Primary Keys section
            if table_info.get('primary_key'):
                markdown_output.append('#### Primary Keys\n')
                pk_list = ', '.join([f'`{pk}`' for pk in table_info['primary_key']])
                markdown_output.append(f'**Primary Key(s):** {pk_list}\n')
        
            # Foreign Keys section
            if table_info.get('foreign_keys'):
                markdown_output.append('#### Foreign Keys\n')
                markdown_output.append('| Local Column(s) | Referenced Table | Referenced Column(s) |')
                markdown_output.append('|---|---|---|')
            
                for fk in table_info['foreign_keys']:
                    local_cols = ', '.join([f'`{col}`' for col in fk['constrained_columns']])
                    ref_table = f'`{fk["referred_table"]}`'
                    ref_cols = ', '.join([f'`{col}`' for col in fk['referred_columns']])
                
                    markdown_output.append(f'| {local_cols} | {ref_table} | {ref_cols} |')
            
                markdown_output.append('')
        
            markdown_output.append('---\n')
    
        return '\n'.join(markdown_output)

    def _format_query_result_to_markdown(self, columns: list, data: list) -> str:
        if not columns:
            return 'No column names available'
            
        if not data:
            return f'Query result is empty\n\n**Column names:** {", ".join(columns)}'
        
        markdown_lines = []
        
        header = '| ' + ' | '.join(columns) + ' |'
        markdown_lines.append(header)
        
        separator = '|' + '|'.join([' --- ' for _ in columns]) + '|'
        markdown_lines.append(separator)
        
        for row in data:
            formatted_row = []
            for value in row:
                if value is None:
                    formatted_row.append('NULL')
                else:
                    cell_value = str(value).replace('|', '\\|')  # 转义管道符
                    formatted_row.append(cell_value)
            
            row_str = '| ' + ' | '.join(formatted_row) + ' |'
            markdown_lines.append(row_str)
        
        result_summary = f'\n**Query result summary:** {len(data)} rows, {len(columns)} columns'
        markdown_lines.append(result_summary)
        
        result = '\n'.join(markdown_lines)
        return result
