#!/usr/bin/env python3
"""
Oracle Publisher Service (Improved)
Fetches data from multiple sources, aggregates it, and posts updates to the blockchain
Includes retry logic, circuit breaker, and comprehensive error handling
"""

import os
import asyncio
import logging
import time
from typing import List, Optional
from decimal import Decimal
from dataclasses import dataclass, field
from enum import Enum
from web3 import Web3
from web3.middleware import geth_poa_middleware
from eth_account import Account
import requests
from prometheus_client import Counter, Gauge, Histogram, start_http_server
import signal
import sys

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Prometheus metrics
updates_sent = Counter('oracle_updates_sent_total', 'Total number of oracle updates sent')
update_errors = Counter('oracle_update_errors_total', 'Total number of oracle update errors')
current_price = Gauge('oracle_current_price', 'Current oracle price')
price_deviation = Gauge('oracle_price_deviation', 'Price deviation from last update')
circuit_breaker_state = Gauge('oracle_circuit_breaker_state', 'Circuit breaker state (0=closed, 1=open, 2=half-open)')
retry_attempts = Counter('oracle_retry_attempts_total', 'Total number of retry attempts')
request_duration = Histogram('oracle_request_duration_seconds', 'Request duration in seconds', ['operation'])


class CircuitBreakerState(Enum):
    CLOSED = 0  # Normal operation
    OPEN = 1    # Failing, reject requests
    HALF_OPEN = 2  # Testing if service recovered


@dataclass
class DataSource:
    """Data source configuration"""
    name: str
    url: str
    parser: str  # JSON path or function name
    weight: float = 1.0
    timeout: int = 5
    retries: int = 3


@dataclass
class OracleConfig:
    """Oracle configuration"""
    rpc_url: str
    aggregator_address: str
    private_key: str
    heartbeat: int = 60  # seconds
    deviation_threshold: float = 0.5  # percentage
    data_sources: List[DataSource] = field(default_factory=list)
    gas_price: Optional[int] = None
    gas_limit: int = 100000
    max_priority_fee: Optional[int] = None
    circuit_breaker_failure_threshold: int = 5
    circuit_breaker_timeout: int = 60  # seconds
    max_retries: int = 3
    retry_backoff_factor: float = 2.0
    health_check_port: int = 8080


class CircuitBreaker:
    """Circuit breaker pattern implementation"""
    
    def __init__(self, failure_threshold: int = 5, timeout: int = 60):
        self.failure_threshold = failure_threshold
        self.timeout = timeout
        self.failure_count = 0
        self.last_failure_time: Optional[float] = None
        self.state = CircuitBreakerState.CLOSED
    
    def call(self, func, *args, **kwargs):
        """Execute function with circuit breaker protection"""
        if self.state == CircuitBreakerState.OPEN:
            if self.last_failure_time is not None and time.time() - self.last_failure_time > self.timeout:
                self.state = CircuitBreakerState.HALF_OPEN
                logger.info("Circuit breaker: Moving to HALF_OPEN state")
            else:
                raise Exception("Circuit breaker is OPEN")
        
        try:
            result = func(*args, **kwargs)
            if self.state == CircuitBreakerState.HALF_OPEN:
                self.state = CircuitBreakerState.CLOSED
                self.failure_count = 0
                logger.info("Circuit breaker: Moving to CLOSED state")
            circuit_breaker_state.set(self.state.value)
            return result
        except Exception as e:
            self.failure_count += 1
            self.last_failure_time = time.time()
            
            if self.failure_count >= self.failure_threshold:
                self.state = CircuitBreakerState.OPEN
                logger.error(f"Circuit breaker: Moving to OPEN state after {self.failure_count} failures")
            
            circuit_breaker_state.set(self.state.value)
            raise e


class OraclePublisher:
    """Oracle publisher service with retry logic and circuit breaker"""
    
    def __init__(self, config: OracleConfig):
        self.config = config
        self.w3 = Web3(Web3.HTTPProvider(config.rpc_url))
        self.w3.middleware_onion.inject(geth_poa_middleware, layer=0)
        self.account = Account.from_key(config.private_key)
        self.last_price: Optional[Decimal] = None
        self.last_update_time: float = 0
        self.aggregator_abi = self._load_aggregator_abi()
        self.aggregator_contract = self.w3.eth.contract(
            address=Web3.to_checksum_address(config.aggregator_address),
            abi=self.aggregator_abi
        )
        self.circuit_breaker = CircuitBreaker(
            failure_threshold=config.circuit_breaker_failure_threshold,
            timeout=config.circuit_breaker_timeout
        )
        self.shutdown_event = asyncio.Event()
        
    def _load_aggregator_abi(self) -> List[dict]:
        """Load aggregator ABI"""
        return [
            {
                "inputs": [{"internalType": "uint256", "name": "answer", "type": "uint256"}],
                "name": "updateAnswer",
                "outputs": [],
                "stateMutability": "nonpayable",
                "type": "function"
            },
            {
                "inputs": [],
                "name": "latestAnswer",
                "outputs": [{"internalType": "int256", "name": "", "type": "int256"}],
                "stateMutability": "view",
                "type": "function"
            },
            {
                "inputs": [],
                "name": "latestRound",
                "outputs": [{"internalType": "uint256", "name": "", "type": "uint256"}],
                "stateMutability": "view",
                "type": "function"
            }
        ]
    
    async def fetch_price_from_source(self, source: DataSource) -> Optional[Decimal]:
        """Fetch price from a single source with retry logic"""
        for attempt in range(source.retries):
            try:
                with request_duration.labels(operation=f"fetch_{source.name}").time():
                    response = requests.get(
                        source.url,
                        timeout=source.timeout,
                        headers={'User-Agent': 'OraclePublisher/1.0'}
                    )
                    response.raise_for_status()
                    data = response.json()
                    
                    # Parse price from JSON (simplified - should use jsonpath in production)
                    parser = source.parser
                    if '.' in parser:
                        keys = parser.split('.')
                        value = data
                        for key in keys:
                            value = value[key]
                    else:
                        value = data.get(parser, data.get('price', data.get('data', {}).get('price')))
                    
                    return Decimal(str(value))
            except Exception as e:
                retry_attempts.inc()
                if attempt < source.retries - 1:
                    wait_time = self.config.retry_backoff_factor ** attempt
                    logger.warning(f"Failed to fetch from {source.name} (attempt {attempt + 1}/{source.retries}): {e}. Retrying in {wait_time}s")
                    await asyncio.sleep(wait_time)
                else:
                    logger.error(f"Failed to fetch from {source.name} after {source.retries} attempts: {e}")
                    return None
        return None
    
    async def fetch_aggregated_price(self) -> Optional[Decimal]:
        """Fetch and aggregate prices from all sources"""
        prices = []
        weights = []
        
        # Fetch prices from all sources in parallel
        tasks = [self.fetch_price_from_source(source) for source in self.config.data_sources]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error(f"Error fetching from {self.config.data_sources[i].name}: {result}")
                continue
            if result is not None:
                prices.append(result)
                weights.append(self.config.data_sources[i].weight)
        
        if not prices:
            logger.error("No prices fetched from any source")
            return None
        
        # Calculate weighted median
        sorted_prices = sorted(prices)
        n = len(sorted_prices)
        
        if n % 2 == 0:
            median = (sorted_prices[n // 2 - 1] + sorted_prices[n // 2]) / 2
        else:
            median = sorted_prices[n // 2]
        
        logger.info(f"Aggregated price (median): {median}")
        return median
    
    def should_update(self, new_price: Decimal) -> bool:
        """Check if price should be updated"""
        current_time = time.time()
        
        # Check heartbeat
        if current_time - self.last_update_time >= self.config.heartbeat:
            logger.info("Update required: heartbeat elapsed")
            return True
        
        # Check deviation threshold
        if self.last_price is None:
            return True
        
        deviation = abs((new_price - self.last_price) / self.last_price * 100)
        price_deviation.set(float(deviation))
        
        if deviation >= self.config.deviation_threshold:
            logger.info(f"Update required: deviation {deviation:.2f}% >= {self.config.deviation_threshold}%")
            return True
        
        logger.debug(f"Update not required: deviation {deviation:.2f}% < {self.config.deviation_threshold}%")
        return False
    
    async def update_oracle(self, price: Decimal) -> bool:
        """Update oracle on blockchain with retry logic"""
        for attempt in range(self.config.max_retries):
            try:
                # Use circuit breaker
                return self.circuit_breaker.call(self._update_oracle_internal, price)
            except Exception as e:
                if attempt < self.config.max_retries - 1:
                    wait_time = self.config.retry_backoff_factor ** attempt
                    logger.warning(f"Failed to update oracle (attempt {attempt + 1}/{self.config.max_retries}): {e}. Retrying in {wait_time}s")
                    retry_attempts.inc()
                    await asyncio.sleep(wait_time)
                else:
                    logger.error(f"Failed to update oracle after {self.config.max_retries} attempts: {e}")
                    update_errors.inc()
                    return False
        return False
    
    def _update_oracle_internal(self, price: Decimal) -> bool:
        """Internal method to update oracle (called by circuit breaker)"""
        # Convert price to uint256 (assuming 8 decimals)
        price_int = int(price * Decimal(10**8))
        
        # Build transaction
        nonce = self.w3.eth.get_transaction_count(self.account.address)
        
        # Get gas price
        if self.config.gas_price:
            gas_price = self.config.gas_price
        else:
            gas_price = self.w3.eth.gas_price
        
        # Build transaction
        transaction = self.aggregator_contract.functions.updateAnswer(price_int).build_transaction({
            'from': self.account.address,
            'nonce': nonce,
            'gas': self.config.gas_limit,
            'gasPrice': gas_price,
            'chainId': 138
        })
        
        # Sign transaction
        signed_txn = self.account.sign_transaction(transaction)
        
        # Send transaction
        tx_hash = self.w3.eth.send_raw_transaction(signed_txn.rawTransaction)
        logger.info(f"Transaction sent: {tx_hash.hex()}")
        
        # Wait for receipt
        receipt = self.w3.eth.wait_for_transaction_receipt(tx_hash, timeout=120)
        if receipt.status == 1:
            logger.info(f"Transaction confirmed: {tx_hash.hex()}")
            updates_sent.inc()
            current_price.set(float(price))
            self.last_price = price
            self.last_update_time = time.time()
            return True
        else:
            logger.error(f"Transaction failed: {tx_hash.hex()}")
            update_errors.inc()
            return False
    
    async def health_check(self):
        """Health check endpoint"""
        from aiohttp import web
        
        async def handle_health(request):
            health_status = {
                'status': 'healthy',
                'circuit_breaker_state': self.circuit_breaker.state.name,
                'last_update_time': self.last_update_time,
                'last_price': str(self.last_price) if self.last_price else None,
                'rpc_connected': self.w3.is_connected()
            }
            return web.json_response(health_status)
        
        app = web.Application()
        app.router.add_get('/health', handle_health)
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, '0.0.0.0', self.config.health_check_port)
        await site.start()
        logger.info(f"Health check server started on port {self.config.health_check_port}")
    
    async def run(self):
        """Main loop"""
        logger.info("Starting oracle publisher service")
        logger.info(f"RPC URL: {self.config.rpc_url}")
        logger.info(f"Aggregator address: {self.config.aggregator_address}")
        logger.info(f"Heartbeat: {self.config.heartbeat}s")
        logger.info(f"Deviation threshold: {self.config.deviation_threshold}%")
        
        # Start Prometheus metrics server
        start_http_server(8000)
        logger.info("Prometheus metrics server started on port 8000")
        
        # Start health check server
        health_task = asyncio.create_task(self.health_check())
        
        while not self.shutdown_event.is_set():
            try:
                # Fetch aggregated price
                price = await self.fetch_aggregated_price()
                
                if price is None:
                    logger.warning("Failed to fetch price, retrying in 10 seconds")
                    await asyncio.sleep(10)
                    continue
                
                # Check if update is needed
                if self.should_update(price):
                    # Update oracle
                    success = await self.update_oracle(price)
                    if success:
                        logger.info(f"Oracle updated successfully: {price}")
                    else:
                        logger.error("Failed to update oracle")
                else:
                    logger.debug(f"Update not required: {price}")
                
                # Wait before next iteration
                await asyncio.sleep(10)
                
            except asyncio.CancelledError:
                logger.info("Shutting down oracle publisher service")
                break
            except Exception as e:
                logger.error(f"Error in main loop: {e}")
                await asyncio.sleep(10)
        
        health_task.cancel()


def load_config() -> OracleConfig:
    """Load configuration from environment variables"""
    return OracleConfig(
        rpc_url=os.getenv("RPC_URL", "http://localhost:8545"),
        aggregator_address=os.getenv("AGGREGATOR_ADDRESS", ""),
        private_key=os.getenv("PRIVATE_KEY", ""),
        heartbeat=int(os.getenv("HEARTBEAT", "60")),
        deviation_threshold=float(os.getenv("DEVIATION_THRESHOLD", "0.5")),
        data_sources=[
            DataSource(
                name="source1",
                url=os.getenv("DATA_SOURCE_1_URL", ""),
                parser=os.getenv("DATA_SOURCE_1_PARSER", "price"),
                weight=1.0,
                timeout=int(os.getenv("DATA_SOURCE_1_TIMEOUT", "5")),
                retries=int(os.getenv("DATA_SOURCE_1_RETRIES", "3"))
            ),
            DataSource(
                name="source2",
                url=os.getenv("DATA_SOURCE_2_URL", ""),
                parser=os.getenv("DATA_SOURCE_2_PARSER", "price"),
                weight=1.0,
                timeout=int(os.getenv("DATA_SOURCE_2_TIMEOUT", "5")),
                retries=int(os.getenv("DATA_SOURCE_2_RETRIES", "3"))
            ),
        ],
        gas_limit=int(os.getenv("GAS_LIMIT", "100000")),
        gas_price=int(os.getenv("GAS_PRICE", "0")) or None,
        max_priority_fee=int(os.getenv("MAX_PRIORITY_FEE", "0")) or None,
        circuit_breaker_failure_threshold=int(os.getenv("CIRCUIT_BREAKER_FAILURE_THRESHOLD", "5")),
        circuit_breaker_timeout=int(os.getenv("CIRCUIT_BREAKER_TIMEOUT", "60")),
        max_retries=int(os.getenv("MAX_RETRIES", "3")),
        retry_backoff_factor=float(os.getenv("RETRY_BACKOFF_FACTOR", "2.0")),
        health_check_port=int(os.getenv("HEALTH_CHECK_PORT", "8080"))
    )


async def main():
    """Main entry point"""
    config = load_config()
    
    if not config.aggregator_address:
        logger.error("AGGREGATOR_ADDRESS environment variable not set")
        return
    
    if not config.private_key:
        logger.error("PRIVATE_KEY environment variable not set")
        return
    
    publisher = OraclePublisher(config)
    
    # Handle shutdown signals
    def signal_handler(sig, frame):
        logger.info("Received shutdown signal")
        publisher.shutdown_event.set()
    
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    
    try:
        await publisher.run()
    except KeyboardInterrupt:
        logger.info("Shutting down oracle publisher service")
    finally:
        logger.info("Oracle publisher service stopped")


if __name__ == "__main__":
    asyncio.run(main())

