from ..base import ModelService
from .config import IMAGE_GENERATION_PRICE, get_config
from typing import Dict, Any, Tuple, Generator, Union
import logging

logger = logging.getLogger(__name__)

class NanoBananaService(ModelService):
    def get_api_config(self) -> Tuple[str, str]:
        config = get_config()
        return config['api_url'], config['api_key']

    def check_balance(self, balance: float) -> Tuple[bool, float, str]:
        if balance < IMAGE_GENERATION_PRICE:
            return False, IMAGE_GENERATION_PRICE, f'余额不足。当前余额: {balance}, 需要: {IMAGE_GENERATION_PRICE}'
        return True, IMAGE_GENERATION_PRICE, ""

    def calculate_cost(self, usage: Dict[str, Any] = None, stream: bool = False) -> float:
        # 固定按次计费
        return IMAGE_GENERATION_PRICE

    def prepare_payload(self, data: Dict[str, Any]) -> Dict[str, Any]:
        return {k: v for k, v in data.items() if k not in ['user']}

    def handle_response(self, response, stream: bool) -> Union[Dict[str, Any], Generator]:
        pass