from ..base import ModelService
from .config import TOKEN_PRICE_INPUT, TOKEN_PRICE_OUTPUT, MIN_BALANCE, get_config
from typing import Dict, Any, Tuple, Generator, Union
import json
import logging

logger = logging.getLogger(__name__)

class DeepSeekService(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 < MIN_BALANCE:
            return False, MIN_BALANCE, f'余额不足。当前余额: {balance}, 需要至少: {MIN_BALANCE}'
        return True, 0.0, ""

    def calculate_cost(self, usage: Dict[str, Any] = None, stream: bool = False) -> float:
        if not usage:
            return 0.0
        
        prompt_tokens = usage.get('prompt_tokens', 0)
        completion_tokens = usage.get('completion_tokens', 0)
        
        cost = (prompt_tokens * TOKEN_PRICE_INPUT) + (completion_tokens * TOKEN_PRICE_OUTPUT)
        return cost

    def prepare_payload(self, data: Dict[str, Any]) -> Dict[str, Any]:
        payload = {k: v for k, v in data.items() if k not in ['user']}
        # 强制 DeepSeek 返回 usage
        if data.get('stream', False):
            if 'stream_options' not in payload:
                payload['stream_options'] = {"include_usage": True}
        return payload

    def handle_response(self, response, stream: bool) -> Union[Dict[str, Any], Generator]:
        # Response handling is mostly done in the route handler for stream/json split
        # but we can provide helper methods to parse usage
        pass
    
    @staticmethod
    def parse_stream_usage(chunk_text: str) -> Dict[str, Any]:
        """解析流式响应中的 usage"""
        try:
            for line in chunk_text.split('\n'):
                if line.startswith('data: ') and line != 'data: [DONE]':
                    json_str = line[6:]
                    data_obj = json.loads(json_str)
                    if 'usage' in data_obj and data_obj['usage']:
                        return data_obj['usage']
        except Exception:
            pass
        return None