import os
import logging
from typing import List, Dict, Any, AsyncGenerator, Optional
import openai
from src.services.provider.base import (
    BaseLLMService, 
    ServiceResponse, 
    TaskStatus, 
    GenerationResult
)

logger = logging.getLogger(__name__)

class OpenAIService(BaseLLMService):
    """ 服务 for OpenAI-compatible LLM providers (replacing LiteLLM).
    Can be used with Moonshot, DeepSeek, DashScope, etc.
    """
    def __init__(self, model_name: str, base_url: Optional[str] = None, api_key: Optional[str] = None, **kwargs):
        super().__init__(model_name, api_key=api_key, **kwargs)
        self.base_url = base_url
        
        if not self.api_key:
            # Fallback to default env var if not injected
            self.api_key = os.environ.get("OPENAI_API_KEY")
        
        if not self.api_key:
            logger.warning(f"API Key not found for {model_name}.")

        self.client = openai.AsyncClient(
            api_key=self.api_key,
            base_url=self.base_url
        )
        logger.info(f"Initialized OpenAIService for {model_name} with base_url={base_url}")

    async def call(self, 
             messages: List[Dict[str, Any]], 
             stream: bool = True,
             temperature: float = 0.7,
             top_p: float = 0.8,
             **kwargs) -> ServiceResponse:
        
        try:
            if stream:
                # streaming, we return the generator directly via a helper if needed,
                # but BaseLLMService signature expects ServiceResponse.
                # However, chat.py handles streaming differently (calling stream() method if available).
                # We'll implement _call_streaming and _call_non_streaming.
                # But wait, BaseLLMService.call returns ServiceResponse. 
                # If stream=True, we usually return a generator in 'results' or handle it differently.
                # Actually, standard pattern in this repo seems to be:
                # call() -> ServiceResponse (non-streaming)
                # stream() -> AsyncGenerator (streaming)
                # But BaseLLMService defines 'call' with 'stream' param.
                
                # Let's implement stream() separately and use call() for non-streaming mostly.
                # If stream=True is passed to call(), we can return a generator?
                # 服务Response fields are static.
                # We will support stream() method explicitly.
                pass

            response = await self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                stream=False,
                **kwargs
            )
            
            content = response.choices[0].message.content
            usage = dict(response.usage) if response.usage else {}
            
            return ServiceResponse(
                status=TaskStatus.SUCCEEDED,
                results=[GenerationResult(content=content, usage=usage)],
                meta={"finish_reason": response.choices[0].finish_reason}
            )

        except Exception as e:
            logger.error(f"Error in OpenAIService call: {e}")
            return ServiceResponse(
                status=TaskStatus.FAILED,
                error=str(e)
            )

    async def stream(self, 
             messages: List[Dict[str, Any]], 
             temperature: float = 0.7,
             top_p: float = 0.8,
             **kwargs) -> AsyncGenerator[str, None]:
        """ 流ing implementation compatible with FastAPI StreamingResponse"""
        try:
            stream = await self.client.chat.completions.create(
                model=self.model_name,
                messages=messages,
                temperature=temperature,
                top_p=top_p,
                stream=True,
                **kwargs
            )
            
            async for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
                    
        except Exception as e:
            logger.error(f"Error in OpenAIService stream: {e}")
            yield f"Error: {str(e)}"
