#!/usr/bin/env python3
"""
Async PostgreSQL gRPC Client
High-performance async PostgreSQL client using grpc.aio

Performance Benefits:
- True async I/O without GIL blocking
- Concurrent query execution
- Memory-efficient connection pooling
"""

from typing import List, Dict, Optional, Any, TYPE_CHECKING
from google.protobuf.struct_pb2 import Struct, Value
from .async_base_client import AsyncBaseGRPCClient
from .proto import postgres_service_pb2, postgres_service_pb2_grpc

if TYPE_CHECKING:
    from .consul_client import ConsulRegistry


class AsyncPostgresClient(AsyncBaseGRPCClient):
    """Async PostgreSQL gRPC client for high-performance concurrent operations."""

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[int] = None,
        user_id: Optional[str] = None,
        lazy_connect: bool = True,
        enable_compression: bool = True,
        enable_retry: bool = True,
        consul_registry: Optional['ConsulRegistry'] = None,
        service_name_override: Optional[str] = None
    ):
        """
        Initialize async PostgreSQL client.

        Args:
            host: Service address (optional, will use Consul discovery if not provided)
            port: Service port (optional, will use Consul discovery if not provided)
            user_id: User ID
            lazy_connect: Lazy connection (default: True)
            enable_compression: Enable compression (default: True)
            enable_retry: Enable retry (default: True)
            consul_registry: ConsulRegistry instance for service discovery (optional)
            service_name_override: Override service name for Consul lookup (optional)
        """
        super().__init__(
            host=host,
            port=port,
            user_id=user_id,
            lazy_connect=lazy_connect,
            enable_compression=enable_compression,
            enable_retry=enable_retry,
            consul_registry=consul_registry,
            service_name_override=service_name_override
        )

    def _create_stub(self):
        """Create PostgreSQL service stub."""
        return postgres_service_pb2_grpc.PostgresServiceStub(self.channel)

    def service_name(self) -> str:
        return "PostgreSQL"

    def default_port(self) -> int:
        return 50061

    # ============================================
    # Health Check
    # ============================================

    async def health_check(self, detailed: bool = True) -> Optional[Dict]:
        """Health check."""
        try:
            await self._ensure_connected()
            request = postgres_service_pb2.HealthCheckRequest(detailed=detailed)
            response = await self.stub.HealthCheck(request)

            return {
                'status': response.status,
                'healthy': response.healthy,
                'version': response.version,
                'details': dict(response.details) if response.details else {}
            }

        except Exception as e:
            return self.handle_error(e, "Health check")

    # ============================================
    # Query Operations
    # ============================================

    async def query(self, sql: str, params: Optional[List[Any]] = None, schema: str = 'public') -> Optional[List[Dict]]:
        """Execute SELECT query.

        Args:
            sql: SQL query statement
            params: Query parameters
            schema: Database schema (default: public)

        Returns:
            List of result rows as dictionaries
        """
        try:
            await self._ensure_connected()

            proto_params = []
            if params:
                for param in params:
                    proto_params.append(self._python_to_proto_value(param))

            request = postgres_service_pb2.QueryRequest(
                sql=sql,
                params=proto_params,
                schema=schema
            )

            response = await self.stub.Query(request)

            if response.metadata.success:
                rows = [self._proto_struct_to_dict(row) for row in response.rows]
                return rows
            return None

        except Exception as e:
            return self.handle_error(e, "Query")

    async def query_row(self, sql: str, params: Optional[List[Any]] = None, schema: str = 'public') -> Optional[Dict]:
        """Execute single row query.

        Args:
            sql: SQL query statement
            params: Query parameters
            schema: Database schema

        Returns:
            Single row as dictionary or None
        """
        try:
            await self._ensure_connected()

            proto_params = []
            if params:
                for param in params:
                    proto_params.append(self._python_to_proto_value(param))

            request = postgres_service_pb2.QueryRowRequest(
                sql=sql,
                params=proto_params,
                schema=schema
            )

            response = await self.stub.QueryRow(request)

            if response.metadata.success and response.found:
                return self._proto_struct_to_dict(response.row)
            return None

        except Exception as e:
            return self.handle_error(e, "Query row")

    async def execute(self, sql: str, params: Optional[List[Any]] = None, schema: str = 'public') -> Optional[int]:
        """Execute INSERT/UPDATE/DELETE statement.

        Args:
            sql: SQL statement
            params: Statement parameters
            schema: Database schema

        Returns:
            Number of rows affected
        """
        try:
            await self._ensure_connected()

            proto_params = []
            if params:
                for param in params:
                    proto_params.append(self._python_to_proto_value(param))

            request = postgres_service_pb2.ExecuteRequest(
                sql=sql,
                params=proto_params,
                schema=schema
            )

            response = await self.stub.Execute(request)

            if response.metadata.success:
                return response.rows_affected
            return None

        except Exception as e:
            return self.handle_error(e, "Execute")

    async def execute_batch(self, operations: List[Dict[str, Any]], schema: str = 'public') -> Optional[Dict]:
        """Execute batch operations.

        Args:
            operations: List of {'sql': str, 'params': List} dictionaries
            schema: Database schema

        Returns:
            Batch execution results
        """
        try:
            await self._ensure_connected()

            batch_ops = []
            for op in operations:
                proto_params = []
                if 'params' in op and op['params']:
                    for param in op['params']:
                        proto_params.append(self._python_to_proto_value(param))

                batch_ops.append(postgres_service_pb2.BatchOperation(
                    sql=op['sql'],
                    params=proto_params
                ))

            request = postgres_service_pb2.ExecuteBatchRequest(
                operations=batch_ops,
                schema=schema
            )

            response = await self.stub.ExecuteBatch(request)

            if response.metadata.success:
                return {
                    'total_rows_affected': response.total_rows_affected,
                    'results': [{'rows_affected': r.rows_affected, 'error': r.error} for r in response.results]
                }
            return None

        except Exception as e:
            return self.handle_error(e, "Execute batch")

    # ============================================
    # Query Builder Operations
    # ============================================

    async def select_from(
        self,
        table: str,
        columns: Optional[List[str]] = None,
        where: Optional[List[Dict]] = None,
        order_by: Optional[List[str]] = None,
        limit: int = 0,
        offset: int = 0,
        schema: str = 'public'
    ) -> Optional[List[Dict]]:
        """Query builder style SELECT.

        Args:
            table: Table name
            columns: Columns to select (default: all)
            where: WHERE conditions as list of dicts
            order_by: ORDER BY clauses
            limit: LIMIT value
            offset: OFFSET value
            schema: Database schema

        Returns:
            List of result rows
        """
        try:
            await self._ensure_connected()

            where_clauses = []
            if where:
                for w in where:
                    where_clauses.append(postgres_service_pb2.WhereClause(
                        column=w.get('column', ''),
                        operator=w.get('operator', '='),
                        value=self._python_to_proto_value(w.get('value'))
                    ))

            request = postgres_service_pb2.SelectFromRequest(
                table=table,
                columns=columns or [],
                where=where_clauses,
                order_by=order_by or [],
                limit=limit,
                offset=offset,
                schema=schema
            )

            response = await self.stub.SelectFrom(request)

            if response.metadata.success:
                return [self._proto_struct_to_dict(row) for row in response.rows]
            return None

        except Exception as e:
            return self.handle_error(e, "Select from")

    async def insert_into(
        self,
        table: str,
        rows: List[Dict],
        returning: bool = False,
        schema: str = 'public'
    ) -> Optional[int]:
        """Insert rows into table.

        Args:
            table: Table name
            rows: List of row dictionaries to insert
            returning: Return inserted rows
            schema: Database schema

        Returns:
            Number of rows inserted
        """
        try:
            await self._ensure_connected()

            proto_rows = []
            for row in rows:
                struct = Struct()
                struct.update(row)
                proto_rows.append(struct)

            request = postgres_service_pb2.InsertIntoRequest(
                table=table,
                rows=proto_rows,
                returning=returning,
                schema=schema
            )

            response = await self.stub.InsertInto(request)

            if response.metadata.success:
                return response.rows_inserted
            return None

        except Exception as e:
            return self.handle_error(e, "Insert into")

    # ============================================
    # Table Operations
    # ============================================

    async def list_tables(self, schema: str = 'public') -> List[str]:
        """List all tables in schema.

        Args:
            schema: Database schema

        Returns:
            List of table names
        """
        try:
            await self._ensure_connected()

            request = postgres_service_pb2.ListTablesRequest(schema=schema)
            response = await self.stub.ListTables(request)

            if response.metadata.success:
                return list(response.tables)
            return []

        except Exception as e:
            self.handle_error(e, "List tables")
            return []

    async def table_exists(self, table: str, schema: str = 'public') -> bool:
        """Check if table exists.

        Args:
            table: Table name
            schema: Database schema

        Returns:
            True if table exists, False otherwise
        """
        try:
            await self._ensure_connected()

            request = postgres_service_pb2.TableExistsRequest(
                table=table,
                schema=schema
            )
            response = await self.stub.TableExists(request)

            if response.metadata.success:
                return response.exists
            return False

        except Exception as e:
            self.handle_error(e, "Table exists check")
            return False

    # ============================================
    # Statistics
    # ============================================

    async def get_stats(self) -> Optional[Dict]:
        """Get connection pool and database statistics.

        Returns:
            Statistics dictionary
        """
        try:
            await self._ensure_connected()

            request = postgres_service_pb2.GetStatsRequest()
            response = await self.stub.GetStats(request)

            if response.metadata.success:
                return {
                    'pool': {
                        'max_connections': response.pool_stats.max_connections,
                        'open_connections': response.pool_stats.open_connections,
                        'idle_connections': response.pool_stats.idle_connections,
                        'active_connections': response.pool_stats.active_connections,
                        'total_queries': response.pool_stats.total_queries,
                    },
                    'database': {
                        'version': response.db_stats.version,
                    }
                }
            return None

        except Exception as e:
            return self.handle_error(e, "Get stats")

    # ============================================
    # Helper Methods
    # ============================================

    def _python_to_proto_value(self, value: Any) -> Value:
        """Convert Python value to proto Value."""
        proto_value = Value()

        if value is None:
            proto_value.null_value = 0
        elif isinstance(value, bool):
            proto_value.bool_value = value
        elif isinstance(value, (int, float)):
            proto_value.number_value = float(value)
        elif isinstance(value, str):
            proto_value.string_value = value
        elif isinstance(value, list):
            proto_value.list_value.values.extend([self._python_to_proto_value(v) for v in value])
        elif isinstance(value, dict):
            for k, v in value.items():
                proto_value.struct_value.fields[k].CopyFrom(self._python_to_proto_value(v))
        else:
            proto_value.string_value = str(value)

        return proto_value

    # ============================================
    # Concurrent Operations
    # ============================================

    async def query_many_concurrent(self, queries: List[Dict[str, Any]]) -> List[Optional[List[Dict]]]:
        """
        Execute multiple queries concurrently.

        Args:
            queries: List of {'sql': str, 'params': list, 'schema': str} dicts

        Returns:
            List of results for each query
        """
        import asyncio

        async def execute_query(q: Dict) -> Optional[List[Dict]]:
            return await self.query(
                sql=q.get('sql', ''),
                params=q.get('params'),
                schema=q.get('schema', 'public')
            )

        return await asyncio.gather(*[execute_query(q) for q in queries])

    async def execute_many_concurrent(self, statements: List[Dict[str, Any]]) -> List[Optional[int]]:
        """
        Execute multiple statements concurrently.

        Args:
            statements: List of {'sql': str, 'params': list, 'schema': str} dicts

        Returns:
            List of rows affected for each statement
        """
        import asyncio

        async def execute_stmt(s: Dict) -> Optional[int]:
            return await self.execute(
                sql=s.get('sql', ''),
                params=s.get('params'),
                schema=s.get('schema', 'public')
            )

        return await asyncio.gather(*[execute_stmt(s) for s in statements])


# Example usage
if __name__ == '__main__':
    import asyncio

    async def main():
        async with AsyncPostgresClient(host='localhost', port=50061, user_id='test_user') as client:
            # Health check
            health = await client.health_check()
            print(f"Health: {health}")

            # Query
            tables = await client.list_tables()
            print(f"Tables: {tables}")

            # Concurrent queries
            results = await client.query_many_concurrent([
                {'sql': 'SELECT 1 as num'},
                {'sql': 'SELECT 2 as num'},
                {'sql': 'SELECT 3 as num'}
            ])
            print(f"Concurrent results: {results}")

    asyncio.run(main())
