import logging
import os
from typing import Optional, Dict, Any, Union, Sequence, AsyncIterator
import functools
import asyncio

from agentscope.model import OpenAIChatModel
from agentscope.message import Msg
from src.services.provider.registry import ModelRegistry

# Import OpenAI types for mocking stream
try:
    from openai.types.chat import ChatCompletionChunk
    from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta, ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction
except ImportError:
    # Fallback if openai is not installed (unlikely)
    ChatCompletionChunk = Any
    Choice = Any
    ChoiceDelta = Any
    ChoiceDeltaToolCall = Any
    ChoiceDeltaToolCallFunction = Any

logger = logging.getLogger(__name__)


class MockAsyncStream:
    """Mocks an OpenAI async stream from a generator."""
    def __init__(self, generator):
        self.generator = generator
    
    async def __aenter__(self):
        return self.generator
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


def _remove_stream_options_if_stream_false(func):
    """
    Decorator: 
    1. Removes stream_options when stream=False
    2. If stream=True, forces stream=False internally to get full response, 
       then wraps it in a MockAsyncStream to satisfy AgentScope's streaming expectation.
       This prevents JSON parsing errors caused by incomplete stream chunks.
    """
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        # Check if stream is requested
        stream_requested = kwargs.get("stream", False)
        
        # If stream is requested, we intercept it
        if stream_requested:
            # Force stream=False to ensure we get complete JSON
            kwargs["stream"] = False
            # Remove stream_options if present as they conflict with stream=False
            kwargs.pop("stream_options", None)
            
        # Also handle explicit stream=False case (cleanup)
        if "stream" in kwargs and not kwargs["stream"]:
             kwargs.pop("stream_options", None)
             
        # Call original function (now with stream=False)
        # This returns a ChatCompletion object (full response)
        response = await func(*args, **kwargs)
        
        # If stream was requested, we need to wrap the response in a generator
        if stream_requested:
            # Create an async generator that yields the content in format AgentScope expects
            async def response_generator():
                # We need to mimic the chunk structure OpenAI returns
                
                content = response.choices[0].message.content
                tool_calls = response.choices[0].message.tool_calls
                
                # 1. Yield content if present
                if content:
                    yield ChatCompletionChunk(
                        id=response.id,
                        choices=[Choice(
                            delta=ChoiceDelta(content=content, role="assistant"),
                            finish_reason=None,
                            index=0
                        )],
                        created=response.created,
                        model=response.model,
                        object="chat.completion.chunk"
                    )
                
                # 2. Yield tool calls if present
                if tool_calls:
                     # For tool calls, we should yield them properly.
                     # AgentScope might expect them in one chunk or multiple.
                     # Yielding in one chunk is safest for parsing.
                     
                     delta_tool_calls = []
                     for i, tc in enumerate(tool_calls):
                         logger.debug(f"Mocking tool call stream for {tc.function.name}")
                         # Use explicit ChoiceDeltaToolCallFunction to ensure Pydantic validation works
                         function_obj = ChoiceDeltaToolCallFunction(
                             name=tc.function.name, 
                             arguments=tc.function.arguments
                         )
                         delta_tool_calls.append(ChoiceDeltaToolCall(
                            index=i,
                            id=tc.id,
                            type=tc.type,
                            function=function_obj
                        ))
                     
                     yield ChatCompletionChunk(
                        id=response.id,
                        choices=[Choice(
                            delta=ChoiceDelta(
                                tool_calls=delta_tool_calls
                            ),
                            finish_reason=None,
                            index=0
                        )],
                        created=response.created,
                        model=response.model,
                        object="chat.completion.chunk"
                    )

                # 3. Yield finish
                yield ChatCompletionChunk(
                    id=response.id,
                    choices=[Choice(
                        delta=ChoiceDelta(),
                        finish_reason=response.choices[0].finish_reason,
                        index=0
                    )],
                    created=response.created,
                    model=response.model,
                    object="chat.completion.chunk"
                )
            
            return MockAsyncStream(response_generator())
            
        return response
    return wrapper

class PixelAgentScopeModel(OpenAIChatModel):
    """ 适配器 that initializes an AgentScope OpenAIChatModel using configuration 
    from Pixel's ModelRegistry.
    
    This simplifies the integration by letting AgentScope handle the actual API calls
    via its native OpenAIChatModel, while we just provide the configuration (api_key, base_url)
    lookup from our centralized registry.
    """
    def __init__(
        self, 
        provider: str,
        model_name: str,
        **kwargs
    ):
        """ 初始化 the model using configuration from the registry.
        
        Args:
            provider: Provider name (e.g., "dashscope", "moonshot", "google")
            model_name: Model name (e.g., "qwen-plus", "kimi-k2.5")
            **kwargs: Additional overrides or arguments for OpenAIChatModel.
        
        Note:
            base_url will be automatically loaded from ModelRegistry configuration.
        """
        # 构建完整的模型 ID
        full_model_id = f"{provider}/{model_name}"
        
        # 1. Fetch config from registry
        service_config = ModelRegistry.get_config(full_model_id)
        
        # 默认 values
        resolved_model_name = model_name
        api_key = None
        base_url = None
        generate_args = {}
        
        if service_config:
            # 提取 core config (model_name, api_key, base_url) from registry
            if isinstance(service_config, dict):
                service_kwargs = service_config.get("kwargs", {})
                provider_name = service_config.get("provider", "openai")
                # Ensure model_name matches what's in the config if not overridden
                # But typically we trust config_name passed in is the ID
            else:
                service_kwargs = service_config.kwargs
                provider_name = service_config.provider
            
            # 1. Model Name
            resolved_model_name = service_kwargs.get("model_name", service_kwargs.get("model", resolved_model_name))

            # 2. Resolve API Key
            api_key = kwargs.get("api_key")
            if not api_key:
                api_key = service_kwargs.get("api_key")
            
            # 3. Base URL
            base_url = service_kwargs.get("base_url")
            if base_url and str(base_url).isupper() and not str(base_url).startswith("http"):
                 resolved_url = os.getenv(base_url)
                 if resolved_url:
                     base_url = resolved_url
            
            # 4. Generate Args
            if "generate_args" in service_kwargs:
                generate_args = service_kwargs["generate_args"]
            
            logger.info(f"Initializing PixelAgentScopeModel for {full_model_id} (model={resolved_model_name})")
        else:
            logger.warning(f"Config for {full_model_id} not found in registry. Using default initialization.")
            # Fallback: try to resolve API key from environment
            # Priority: SCRIPT_*_API_KEY (for script service) > general env vars
            if full_model_id.startswith("qwen") or full_model_id.startswith("dashscope") or (provider and provider == "dashscope"):
                api_key = os.getenv("SCRIPT_DASHSCOPE_API_KEY") or os.getenv("DASHSCOPE_API_KEY")
                base_url = os.getenv("SCRIPT_DASHSCOPE_BASE_URL") or os.getenv("DASHSCOPE_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
            else:
                api_key = os.getenv("SCRIPT_OPENAI_API_KEY") or os.getenv("OPENAI_API_KEY")
                base_url = os.getenv("SCRIPT_OPENAI_BASE_URL") or os.getenv("OPENAI_BASE_URL")
        
        # Construct initialization arguments for OpenAIChatModel
        # 优先级: kwargs (overrides) > registry config > defaults
        
        client_kwargs = {}
        if base_url:
            client_kwargs["base_url"] = base_url
        
        # DashScope 要求 stream=False 时不能设置 stream_options
        # 确保移除任何可能存在的 stream_options
        generate_args.pop("stream_options", None)
            
        init_kwargs = {
            "model_name": resolved_model_name,
            "api_key": api_key,
            "organization": kwargs.pop("organization", None),
            "client_kwargs": client_kwargs,
            "generate_kwargs": generate_args
        }
        
        # 合并 kwargs (overrides)
        if "api_key" in kwargs:
            init_kwargs["api_key"] = kwargs.pop("api_key")
        if "base_url" in kwargs:
            client_kwargs["base_url"] = kwargs.pop("base_url")
        if "model_name" in kwargs:
            init_kwargs["model_name"] = kwargs.pop("model_name")
        if "generate_args" in kwargs:
            init_kwargs["generate_kwargs"] = kwargs.pop("generate_args")
            
        # Add remaining kwargs
        init_kwargs.update(kwargs)
        
        # Remove keys that are None if that causes issues?
        # OpenAIChatModel typically defaults api_key to env var if not passed.
        # But we passed it explicitly (even if None).
        # Let's clean up None values for core params to be safe
        keys_to_clean = ["api_key", "base_url", "organization"]
        for k in keys_to_clean:
            if init_kwargs.get(k) is None:
                init_kwargs.pop(k, None)

        # 纯用户密钥模式：如果没有 api_key，使用占位符避免初始化失败
        # 实际的 api_key 会在调用前通过 set_api_key 方法设置
        self._pending_api_key = None
        if "api_key" not in init_kwargs:
            self._pending_api_key = "placeholder-must-be-set-before-use"
            init_kwargs["api_key"] = self._pending_api_key
            logger.warning(f"PixelAgentScopeModel initialized with placeholder api_key for {full_model_id}. "
                          f"API key must be set via set_api_key() before using this model.")

        super().__init__(**init_kwargs)
        
        # 猴子补丁：包装 client.chat.completions.create 方法
        # 确保当 stream=False 时移除 stream_options
        original_create = self.client.chat.completions.create
        self.client.chat.completions.create = _remove_stream_options_if_stream_false(original_create)

    def set_api_key(self, api_key: str):
        """动态设置 API Key（用于纯用户密钥模式）"""
        import openai
        self.api_key = api_key
        # 重新创建 client 以使用新的 api_key
        client_kwargs = {}
        if hasattr(self, 'client') and hasattr(self.client, '_base_url'):
            client_kwargs['base_url'] = str(self.client._base_url)
        self.client = openai.AsyncClient(api_key=api_key, **client_kwargs)
        # 重新应用猴子补丁
        original_create = self.client.chat.completions.create
        self.client.chat.completions.create = _remove_stream_options_if_stream_false(original_create)
        self._pending_api_key = None
        logger.debug(f"API key updated for model {self.model_name}")

    def _ensure_api_key_set(self):
        """检查 API Key 是否已设置（非占位符）"""
        if self._pending_api_key and self._pending_api_key.startswith("placeholder"):
            raise RuntimeError(
                f"API key not configured for model {self.model_name}. "
                f"Please set SCRIPT_SERVICE_API_KEY environment variable or call set_api_key() before using this model."
            )
