"""
SpyMonk-DB SpanServer (Database Server).

Main gRPC server for SpyMonk-DB.
"""

import logging
import asyncio
from concurrent import futures
import grpc
from pathlib import Path
from typing import Optional

from spymonk_enterprise.network.grpc.gen import spymonk_pb2
from spymonk_enterprise.network.grpc.gen import spymonk_pb2_grpc
from spymonk_enterprise.storage.mvcc import MVCCStore
from spymonk_enterprise.transaction.transaction import TransactionManager, Transaction
from spymonk_enterprise.time.hybrid_clock import HybridLogicalClock, Timestamp
from spymonk_enterprise.sql.executor.executor import QueryExecutor
from spymonk_enterprise.schema.schema import SchemaRegistry

logger = logging.getLogger(__name__)

class SpyMonkDBService(spymonk_pb2_grpc.SpyMonkDBServicer):
    """gRPC service implementation for SpyMonk-DB"""

    def __init__(self, data_dir: str, node_id: str):
        self.data_dir = Path(data_dir)
        self.node_id = node_id
        
        # Initialize core components
        self.clock = HybridLogicalClock(node_id)
        self.clock.start()
        
        self.store = MVCCStore(self.data_dir, self.clock)
        self.txn_manager = TransactionManager(self.clock, self.store)
        self.schema_registry = SchemaRegistry()
        self.executor = QueryExecutor(self.store, self.schema_registry)
        
        logger.info(f"Initialized SpyMonkDBService on node {node_id}")

    def _to_pb_timestamp(self, ts: Timestamp) -> spymonk_pb2.Timestamp:
        return spymonk_pb2.Timestamp(
            physical=ts.physical,
            logical=ts.logical,
            uncertainty_ns=ts.uncertainty_ns
        )

    def _from_pb_timestamp(self, ts_pb: spymonk_pb2.Timestamp) -> Timestamp:
        return Timestamp(
            physical=ts_pb.physical,
            logical=ts_pb.logical,
            uncertainty_ns=ts_pb.uncertainty_ns
        )

    # Transaction operations

    def BeginTransaction(self, request, context):
        read_only = request.read_only
        txn = self.txn_manager.begin(read_only=read_only)
        
        return spymonk_pb2.BeginTransactionResponse(
            transaction_id=txn.txn_id,
            start_timestamp=self._to_pb_timestamp(txn.start_timestamp)
        )

    def CommitTransaction(self, request, context):
        txn_id = request.transaction_id
        txn = self.txn_manager.active_txns.get(txn_id)
        
        if not txn:
            context.abort(grpc.StatusCode.NOT_FOUND, f"Transaction {txn_id} not found")
        
        success = txn.commit()
        return spymonk_pb2.CommitTransactionResponse(
            success=success,
            commit_timestamp=self._to_pb_timestamp(txn.commit_timestamp) if success else None
        )

    def AbortTransaction(self, request, context):
        txn_id = request.transaction_id
        txn = self.txn_manager.active_txns.get(txn_id)
        
        if txn:
            txn.abort()
        
        return spymonk_pb2.AbortTransactionResponse(success=True)

    # Data operations

    def Get(self, request, context):
        txn_id = request.transaction_id
        ts = None
        if request.read_timestamp.physical > 0:
            ts = self._from_pb_timestamp(request.read_timestamp)
            
        if txn_id:
            txn = self.txn_manager.active_txns.get(txn_id)
            if not txn:
                context.abort(grpc.StatusCode.NOT_FOUND, f"Transaction {txn_id} not found")
            value = txn.get(request.key)
        else:
            value = self.store.get(request.key, ts)
            
        return spymonk_pb2.GetResponse(
            value=value or b"",
            found=value is not None
        )

    def Put(self, request, context):
        txn_id = request.transaction_id
        if txn_id:
            txn = self.txn_manager.active_txns.get(txn_id)
            if not txn:
                context.abort(grpc.StatusCode.NOT_FOUND, f"Transaction {txn_id} not found")
            txn.put(request.key, request.value)
            return spymonk_pb2.PutResponse(success=True)
        else:
            ts = self.store.put(request.key, request.value)
            return spymonk_pb2.PutResponse(
                success=True,
                timestamp=self._to_pb_timestamp(ts)
            )

    def Delete(self, request, context):
        txn_id = request.transaction_id
        if txn_id:
            txn = self.txn_manager.active_txns.get(txn_id)
            if not txn:
                context.abort(grpc.StatusCode.NOT_FOUND, f"Transaction {txn_id} not found")
            txn.delete(request.key)
            return spymonk_pb2.DeleteResponse(success=True)
        else:
            ts = self.store.delete(request.key)
            return spymonk_pb2.DeleteResponse(
                success=True,
                timestamp=self._to_pb_timestamp(ts)
            )

    def Scan(self, request, context):
        txn_id = request.transaction_id
        start_key = request.start_key if request.start_key else None
        end_key = request.end_key if request.end_key else None
        ts = None
        if request.read_timestamp.physical > 0:
            ts = self._from_pb_timestamp(request.read_timestamp)
            
        # MVCCStore.scan yields (key, value)
        # We need to yield ScanResponse
        for key, value in self.store.scan(start_key, end_key, ts):
            yield spymonk_pb2.ScanResponse(key=key, value=value)

    # SQL operations

    def ExecuteSQL(self, request, context):
        # Implementation depends on a parser that can turn SQL string into AST
        # For now, let's assume a simplified execution
        from spymonk_enterprise.sql.parser.sql_parser import SQLParser
        
        parser = SQLParser()
        try:
            statement = parser.parse(request.sql)
            
            txn_id = request.transaction_id
            txn = None
            if txn_id:
                txn = self.txn_manager.active_txns.get(txn_id)
                if not txn:
                    context.abort(grpc.StatusCode.NOT_FOUND, f"Transaction {txn_id} not found")
            
            result = self.executor.execute(statement, txn)
            
            rows_pb = []
            for row in result.rows:
                # Convert dict to row with bytes
                cols = {k: str(v).encode() for k, v in row.items()}
                rows_pb.append(spymonk_pb2.QueryRow(columns=cols))
                
            return spymonk_pb2.ExecuteSQLResponse(
                rows=rows_pb,
                affected_rows=result.affected_rows
            )
        except Exception as e:
            logger.error(f"SQL execution error: {e}")
            context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(e))


class SpyMonkServer:
    """Main server class for SpanServer"""

    def __init__(self, host="0.0.0.0", port=50051, data_dir="/tmp/spymonk", node_id="node1"):
        self.host = host
        self.port = port
        self.data_dir = data_dir
        self.node_id = node_id
        self.server = None

    def start(self):
        self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
        
        service = SpyMonkDBService(self.data_dir, self.node_id)
        spymonk_pb2_grpc.add_SpyMonkDBServicer_to_server(service, self.server)
        
        listen_addr = f"{self.host}:{self.port}"
        self.server.add_insecure_port(listen_addr)
        
        self.server.start()
        logger.info(f"SpyMonk-DB Server started on {listen_addr}")
        return self.server

    def stop(self):
        if self.server:
            self.server.stop(0)
            logger.info("SpyMonk-DB Server stopped")
