from typing import Any, List, Optional, Dict, Iterator
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage, SystemMessage, ToolMessage, AIMessageChunk
from langchain_core.outputs import ChatResult, ChatGeneration, ChatGenerationChunk
from app.core.llm.glm_client import glm_client
from app.config import settings

class ChatGLM(BaseChatModel):
    """
    LangChain adapter for ZhipuAI GLM models.
    Wraps the project's existing GLMClient.
    """
    model_name: str = "glm-4.7"
    temperature: float = 0.7
    
    def __init__(self, model: str = None, temperature: float = 0.7, **kwargs):
        super().__init__(**kwargs)
        self.model_name = model or settings.zai_model
        self.temperature = temperature

    @property
    def _llm_type(self) -> str:
        return "chat-glm"

    def _convert_messages(self, messages: List[BaseMessage]) -> List[Dict[str, Any]]:
        glm_messages = []
        for msg in messages:
            role = "user"
            if isinstance(msg, SystemMessage):
                role = "system"
            elif isinstance(msg, AIMessage):
                role = "assistant"
                if msg.tool_calls:
                     # Handle tool calls if necessary, but basic text is priority
                     pass
            elif isinstance(msg, ToolMessage):
                role = "tool"
            
            content = msg.content
            if isinstance(content, str):
                glm_messages.append({"role": role, "content": content})
                
            # Note: Tool calling support might need more complex conversion
            # but for basic chat and text generation this should suffice.
            
        return glm_messages

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        glm_messages = self._convert_messages(messages)
        
        # Use sync call or async call? BaseChatModel._generate is sync.
        # But GLMClient is async-first.
        # We might need to use _agenerate instead or run async in sync.
        # For simplicity, let's use the synchronous client from zai-sdk if possible, 
        # but GLMClient wraps it.
        # Let's check GLMClient again. It has async chat.
        
        # If we are in an async environment (FastAPI), we should implement _agenerate.
        raise NotImplementedError("Use ainvoke or astream for this model")

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        glm_messages = self._convert_messages(messages)
        
        response = await glm_client.chat(
            messages=glm_messages,
            temperature=self.temperature,
            stream=False
        )
        
        # Parse response
        content = response["choices"][0]["message"]["content"]
        
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=content))])

    def _stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        # Sync stream not supported
        raise NotImplementedError("Use astream")

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        glm_messages = self._convert_messages(messages)
        
        response_data = await glm_client.chat(
            messages=glm_messages,
            temperature=self.temperature,
            stream=True
        )
        
        async_gen = response_data["stream"]
        
        async for chunk in async_gen:
            if chunk:
                yield ChatGenerationChunk(message=AIMessageChunk(content=chunk))
