from dataclasses import dataclass
import logging

logger = logging.getLogger(__name__)

class BudgetExceededError(Exception):
    """Raised when a budget limit is exceeded."""
    pass

@dataclass
class BudgetConfig:
    max_queries: int = 5
    max_credits: float = 100.0
    max_schema_calls: int = 3

class BudgetManager:
    def __init__(self, config: BudgetConfig):
        self.config = config
        self.query_count = 0
        self.schema_calls = 0
        self.used_credits = 0.0
        
    def check_can_execute_query(self, estimated_cost: float = 0.0) -> None:
        """
        Check if a new query can be executed within the session budget.
        Raises BudgetExceededError if limits would be breached.
        """
        # Check query count limit
        if self.query_count >= self.config.max_queries:
            raise BudgetExceededError(
                f"Session limit reached: Maximum {self.config.max_queries} queries allowed. "
                "Please search for existing queries or optimize your workflow."
            )
            
        # Check credit limit
        if self.used_credits + estimated_cost > self.config.max_credits:
            remaining = self.config.max_credits - self.used_credits
            raise BudgetExceededError(
                f"Credit limit exceeded: Query costs ~{estimated_cost} credits, "
                f"but only {remaining:.2f} remaining in session budget."
            )
            
    def check_can_access_schema(self) -> None:
        """
        Check if schema access is allowed.
        """
        if self.schema_calls >= self.config.max_schema_calls:
            raise BudgetExceededError(
                f"Schema access limit reached: Max {self.config.max_schema_calls} calls allowed. "
                "Schema lookups are expensive. Please use 'search_public_queries' instead."
            )
            
    def track_execution(self, cost: float) -> None:
        """Record a successful execution."""
        self.query_count += 1
        self.used_credits += cost
        logger.info(f"Budget update: {self.query_count}/{self.config.max_queries} queries, "
                    f"{self.used_credits:.2f}/{self.config.max_credits} credits used.")

    def track_schema_access(self) -> None:
        """Record a schema access."""
        self.schema_calls += 1
        logger.info(f"Budget update: {self.schema_calls}/{self.config.max_schema_calls} schema calls.")

    def sync_usage(self, actual_used: float, actual_limit: float) -> None:
        """
        Sync local budget tracking with the actual usage reported by the Dune API.
        This corrects any drift from our estimations.
        """
        # We only update 'used_credits' based on the API's report for the current period.
        # Note: API usage is global (monthly), while our session budget is local.
        # A simple strategy: If actual usage jumped significantly, we deduct that diff from our session.
        # But for simplicity and safety, let's just log it and perhaps warn if we are near the global limit.
        
        # Actually, for a session manager, we care about "Credits consumed THIS SESSION".
        # Syncing with global usage is hard unless we stored the "start_usage".
        # Let's just update the limit if the API says we have less remaining than we thought.
        
        remaining_global = actual_limit - actual_used
        remaining_session = self.config.max_credits - self.used_credits
        
        if remaining_global < remaining_session:
            # If global remaining is LESS than what we think we have for the session,
            # we must cap our session budget to reality.
            self.config.max_credits = self.used_credits + remaining_global
            logger.info(f"Budget synced: Session limit lowered to match global remaining ({remaining_global:.2f}).")

    def get_status(self) -> dict:
        return {
            "queries": {
                "used": self.query_count,
                "limit": self.config.max_queries,
                "remaining": self.config.max_queries - self.query_count
            },
            "credits": {
                "used": self.used_credits,
                "limit": self.config.max_credits,
                "remaining": self.config.max_credits - self.used_credits
            },
            "schema_calls": {
                "used": self.schema_calls,
                "limit": self.config.max_schema_calls,
                "remaining": self.config.max_schema_calls - self.schema_calls
            }
        }
