#!/usr/bin/env python3
"""
Dragonchain Smart Contract Client

A gRPC client that connects to Dragonchain Prime server to process
smart contract transactions.

Do not modify this file unless you need to customize the client behavior.
Implement your smart contract logic in process.py instead.
"""

import argparse
import json
import logging
import queue
import signal
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Optional

import grpc
import yaml

import remote_sc_pb2 as pb
import remote_sc_pb2_grpc as pb_grpc
from process import ProcessResult, process

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("SmartContract")


# =============================================================================
# Configuration and Client Infrastructure
# Do not modify this file unless you need to customize the client behavior.
# Implement your smart contract logic in process.py instead.
# =============================================================================


@dataclass
class Config:
    """Client configuration loaded from YAML."""

    server_address: str
    smart_contract_id: str
    api_key: str
    use_tls: bool = False
    tls_cert_path: Optional[str] = None
    num_workers: int = 10
    reconnect_delay_seconds: int = 5
    max_reconnect_attempts: int = 0  # 0 = infinite


class SmartContractClient:
    """gRPC client for smart contract execution."""

    def __init__(self, config: Config):
        self.config = config
        self.channel: Optional[grpc.Channel] = None
        self.stub: Optional[pb_grpc.SmartContractServiceStub] = None
        self.running = False
        self.work_queue: queue.Queue = queue.Queue()
        self.response_queue: queue.Queue = queue.Queue()
        self.executor: Optional[ThreadPoolExecutor] = None

    def connect(self) -> bool:
        """Establish connection to the gRPC server."""
        try:
            if self.config.use_tls:
                if not self.config.tls_cert_path:
                    logger.error("TLS enabled but no certificate path provided")
                    return False
                with open(self.config.tls_cert_path, "rb") as f:
                    creds = grpc.ssl_channel_credentials(f.read())
                self.channel = grpc.secure_channel(self.config.server_address, creds)
            else:
                self.channel = grpc.insecure_channel(self.config.server_address)

            self.stub = pb_grpc.SmartContractServiceStub(self.channel)
            logger.info(f"Connected to server at {self.config.server_address}")
            return True
        except Exception as e:
            logger.error(f"Failed to connect: {e}")
            return False

    def close(self):
        """Close the gRPC connection."""
        if self.channel:
            self.channel.close()
            self.channel = None
            self.stub = None

    def _response_generator(self):
        """Generator that yields responses from the response queue."""
        while self.running:
            try:
                response = self.response_queue.get(timeout=1.0)
                if response is None:
                    break
                yield response
            except queue.Empty:
                continue

    def _process_request(self, request: pb.SmartContractRequest):
        """Process a single request and queue the response."""
        logs = ""

        try:
            result = process(
                tx_json=request.transaction_json,
                env_vars=dict(request.env_vars),
                secrets=dict(request.secrets),
            )

            response = pb.SmartContractResponse(
                transaction_id=request.transaction_id,
                output_to_chain=result.output_to_chain,
                logs=logs,
            )

            if result.error:
                response.error = result.error
                logger.error(
                    f"Error processing transaction {request.transaction_id}: {result.error}"
                )
            else:
                response.result_json = json.dumps(result.data) if result.data else "{}"
                logger.info(f"Successfully processed transaction {request.transaction_id}")

        except Exception as e:
            response = pb.SmartContractResponse(
                transaction_id=request.transaction_id,
                error=str(e),
                logs=logs,
            )
            logger.exception(f"Exception processing transaction {request.transaction_id}")

        self.response_queue.put(response)

    def _worker(self):
        """Worker thread that processes requests from the queue."""
        while self.running:
            try:
                request = self.work_queue.get(timeout=1.0)
                if request is None:
                    break
                self._process_request(request)
            except queue.Empty:
                continue

    def run(self) -> bool:
        """Run the client and process incoming requests."""
        if not self.stub:
            logger.error("Not connected to server")
            return False

        self.running = True
        self.executor = ThreadPoolExecutor(max_workers=self.config.num_workers)

        # Start worker threads
        workers = []
        for _ in range(self.config.num_workers):
            future = self.executor.submit(self._worker)
            workers.append(future)

        logger.info(f"Started {self.config.num_workers} worker threads")

        # Create metadata for authentication
        metadata = [
            ("x-api-key", self.config.api_key),
            ("x-smart-contract-id", self.config.smart_contract_id),
        ]

        try:
            # Establish bi-directional stream
            stream = self.stub.Run(self._response_generator(), metadata=metadata)

            logger.info("Stream established, waiting for requests...")

            # Receive and dispatch requests
            for request in stream:
                if not self.running:
                    break
                logger.info(f"Received request: transaction_id={request.transaction_id}")
                self.work_queue.put(request)

            logger.info("Server closed the stream")
            return True

        except grpc.RpcError as e:
            logger.error(f"gRPC error: {e.code()} - {e.details()}")
            return False
        except Exception as e:
            logger.exception(f"Error in run loop: {e}")
            return False
        finally:
            self.running = False

            # Signal workers to stop
            for _ in range(self.config.num_workers):
                self.work_queue.put(None)
            self.response_queue.put(None)

            # Wait for workers to finish
            if self.executor:
                self.executor.shutdown(wait=True)

    def stop(self):
        """Stop the client gracefully."""
        logger.info("Stopping client...")
        self.running = False


def load_config(path: str) -> Config:
    """Load configuration from a YAML file."""
    with open(path, "r") as f:
        data = yaml.safe_load(f)

    # Validate required fields
    required = ["server_address", "smart_contract_id", "api_key"]
    for field in required:
        if field not in data or not data[field]:
            raise ValueError(f"Missing required config field: {field}")

    return Config(
        server_address=data["server_address"],
        smart_contract_id=data["smart_contract_id"],
        api_key=data["api_key"],
        use_tls=data.get("use_tls", False),
        tls_cert_path=data.get("tls_cert_path"),
        num_workers=data.get("num_workers", 10),
        reconnect_delay_seconds=data.get("reconnect_delay_seconds", 5),
        max_reconnect_attempts=data.get("max_reconnect_attempts", 0),
    )


def main():
    parser = argparse.ArgumentParser(description="Dragonchain Smart Contract Client")
    parser.add_argument(
        "--config",
        "-c",
        default="config.yaml",
        help="Path to configuration file",
    )
    args = parser.parse_args()

    # Load configuration
    try:
        config = load_config(args.config)
    except Exception as e:
        logger.error(f"Failed to load config: {e}")
        sys.exit(1)

    # Create client
    client = SmartContractClient(config)

    # Setup signal handling for graceful shutdown
    def signal_handler(signum, frame):
        logger.info(f"Received signal {signum}, shutting down...")
        client.stop()

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # Connection loop with reconnection logic
    attempts = 0
    while True:
        if client.connect():
            attempts = 0
            if not client.run():
                if not client.running:
                    logger.info("Shutdown requested")
                    break

        client.close()

        attempts += 1
        if config.max_reconnect_attempts > 0 and attempts >= config.max_reconnect_attempts:
            logger.error(f"Max reconnection attempts ({config.max_reconnect_attempts}) reached")
            break

        delay = config.reconnect_delay_seconds
        logger.info(f"Reconnecting in {delay} seconds (attempt {attempts})...")
        time.sleep(delay)

    logger.info("Client shut down")


if __name__ == "__main__":
    main()