import logging
import time
from typing import Dict, List, Optional, Any
import httpx

from app.services.dex.base import BaseDexService
from app.services.solana import USDC_TOKEN_ADDRESS

logger = logging.getLogger(__name__)

# Raydium API endpoints
RAYDIUM_API_BASE = "https://api.raydium.io/v2"
PAIRS_ENDPOINT = f"{RAYDIUM_API_BASE}/main/pairs"


class RaydiumDexService(BaseDexService):
    """Service for Raydium DEX price monitoring"""
    
    def __init__(self):
        super().__init__("raydium")
        self.http_client = httpx.AsyncClient(timeout=10.0)
        self.pairs_cache = {}
        self.last_refresh = 0
        self.cache_ttl = 60  # 60 seconds
    
    async def refresh_pairs_cache(self):
        """Refresh the pairs cache if needed"""
        current_time = time.time()
        if current_time - self.last_refresh < self.cache_ttl and self.pairs_cache:
            return
        
        try:
            response = await self.http_client.get(PAIRS_ENDPOINT)
            response.raise_for_status()
            
            pairs_data = response.json()
            
            # Reorganize by token address for faster lookups
            pairs_by_token = {}
            for pair in pairs_data:
                base_token = pair.get("baseMint")
                quote_token = pair.get("quoteMint")
                
                if base_token:
                    if base_token not in pairs_by_token:
                        pairs_by_token[base_token] = []
                    pairs_by_token[base_token].append(pair)
                
                if quote_token:
                    if quote_token not in pairs_by_token:
                        pairs_by_token[quote_token] = []
                    pairs_by_token[quote_token].append(pair)
            
            self.pairs_cache = pairs_by_token
            self.last_refresh = current_time
            
        except Exception as e:
            logger.error(f"Error refreshing Raydium pairs cache: {str(e)}")
    
    async def get_token_price(self, token_address: str, quote_token_address: Optional[str] = None) -> Dict[str, Any]:
        """Get token price from Raydium"""
        if not quote_token_address:
            quote_token_address = USDC_TOKEN_ADDRESS
        
        try:
            await self.refresh_pairs_cache()
            
            # Find all pairs for the token
            token_pairs = self.pairs_cache.get(token_address, [])
            
            # Find a pair with the quote token
            target_pair = None
            for pair in token_pairs:
                if pair.get("baseMint") == token_address and pair.get("quoteMint") == quote_token_address:
                    target_pair = pair
                    price = 1.0 / float(pair.get("price", 0)) if float(pair.get("price", 0)) > 0 else 0
                    break
                elif pair.get("quoteMint") == token_address and pair.get("baseMint") == quote_token_address:
                    target_pair = pair
                    price = float(pair.get("price", 0))
                    break
            
            if target_pair:
                # Calculate liquidity
                amm_id = target_pair.get("ammId")
                liquidity = float(target_pair.get("liquidity", 0))
                
                return {
                    "price": price,
                    "liquidity": liquidity,
                    "timestamp": int(time.time()),
                    "metadata": {
                        "pair_data": target_pair,
                        "amm_id": amm_id
                    }
                }
            else:
                logger.warning(f"No Raydium pair found for {token_address} with quote {quote_token_address}")
                return {
                    "price": 0.0,
                    "liquidity": 0.0,
                    "timestamp": int(time.time()),
                    "metadata": {
                        "error": "No pair found"
                    }
                }
                
        except Exception as e:
            logger.error(f"Error getting Raydium price for {token_address}: {str(e)}")
            return {
                "price": 0.0,
                "liquidity": 0.0,
                "timestamp": int(time.time()),
                "metadata": {
                    "error": str(e)
                }
            }
    
    async def get_token_prices(self, token_addresses: List[str], quote_token_address: Optional[str] = None) -> Dict[str, Dict[str, Any]]:
        """Get prices for multiple tokens from Raydium"""
        if not quote_token_address:
            quote_token_address = USDC_TOKEN_ADDRESS
        
        try:
            await self.refresh_pairs_cache()
            
            result = {}
            timestamp = int(time.time())
            
            for token_address in token_addresses:
                # Find all pairs for the token
                token_pairs = self.pairs_cache.get(token_address, [])
                
                # Find a pair with the quote token
                target_pair = None
                price = 0.0
                
                for pair in token_pairs:
                    if pair.get("baseMint") == token_address and pair.get("quoteMint") == quote_token_address:
                        target_pair = pair
                        price = 1.0 / float(pair.get("price", 0)) if float(pair.get("price", 0)) > 0 else 0
                        break
                    elif pair.get("quoteMint") == token_address and pair.get("baseMint") == quote_token_address:
                        target_pair = pair
                        price = float(pair.get("price", 0))
                        break
                
                if target_pair:
                    # Calculate liquidity
                    amm_id = target_pair.get("ammId")
                    liquidity = float(target_pair.get("liquidity", 0))
                    
                    result[token_address] = {
                        "price": price,
                        "liquidity": liquidity,
                        "timestamp": timestamp,
                        "metadata": {
                            "pair_data": target_pair,
                            "amm_id": amm_id
                        }
                    }
                else:
                    result[token_address] = {
                        "price": 0.0,
                        "liquidity": 0.0,
                        "timestamp": timestamp,
                        "metadata": {
                            "error": "No pair found"
                        }
                    }
            
            return result
                
        except Exception as e:
            logger.error(f"Error getting Raydium prices: {str(e)}")
            timestamp = int(time.time())
            return {
                token_address: {
                    "price": 0.0,
                    "liquidity": 0.0,
                    "timestamp": timestamp,
                    "metadata": {
                        "error": str(e)
                    }
                }
                for token_address in token_addresses
            }