import os
from typing import List, Dict, Any, Optional, Union, Callable
from .base import LLMProvider

class GeminiProvider(LLMProvider):
    provider_name = "gemini"
    
    def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs):
        from google import genai
        
        self.model_name = model_name
        self.api_key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
        if not self.api_key:
            raise ValueError("Gemini/Google API key is required.")
        
        self.client = genai.Client(api_key=self.api_key)

    async def chat(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict[str, Any]]] = None) -> Any:
        from .utils import extract_content
        from google.genai import types
        import json
        
        # Build contents list for the new SDK
        contents = []
        system_instruction = None
        
        for msg in messages:
            role = msg.get("role")
            raw_content = msg.get("content", "")
            tool_calls = msg.get("tool_calls")
            
            if role == "system":
                # Extract text from system message
                if isinstance(raw_content, str):
                    system_instruction = raw_content
                elif isinstance(raw_content, list):
                    text_parts = [p.get("text", "") for p in raw_content if p.get("type") == "text"]
                    system_instruction = " ".join(text_parts)
                continue
            
            parts = []
            
            # 1. Handle regular content (text/images)
            if raw_content:
                text_content, images = extract_content(raw_content)
                if text_content:
                    parts.append(types.Part.from_text(text=text_content))
                
                for img in images:
                    if img["data"] and img["mime_type"]:
                        parts.append(types.Part.from_bytes(
                            data=img["data"],
                            mime_type=img["mime_type"]
                        ))
            
            # 2. Handle assistant tool calls (outgoing history)
            if role == "assistant" and tool_calls:
                for tc in tool_calls:
                    func_data = tc.get("function", {})
                    name = func_data.get("name")
                    args_raw = func_data.get("arguments")
                    
                    if isinstance(args_raw, str):
                        try:
                            args = json.loads(args_raw)
                        except Exception:
                            args = {}
                    else:
                        args = args_raw or {}
                        
                    parts.append(types.Part.from_function_call(
                        name=name,
                        args=args
                    ))
            
            # 3. Handle tool results (incoming history)
            if role == "tool":
                name = msg.get("name")
                # Fallback: recover name from tool_call_id if generated by this provider
                if not name and msg.get("tool_call_id"):
                    tid = msg.get("tool_call_id")
                    if tid.startswith("call_"):
                        name = tid[5:]
                
                tool_content = msg.get("content")
                
                # Parse content to dict for function_response
                if isinstance(tool_content, str):
                    try:
                        resp_dict = json.loads(tool_content)
                        if not isinstance(resp_dict, dict):
                            resp_dict = {"result": resp_dict}
                    except Exception:
                        resp_dict = {"result": tool_content}
                else:
                    resp_dict = tool_content or {}
                
                contents.append(types.Content(
                    role="tool",
                    parts=[types.Part.from_function_response(
                        name=name or "unknown_function", # Avoid empty name
                        response=resp_dict
                    )]
                ))
                continue

            # Add to contents
            if role == "user":
                if parts:
                    contents.append(types.Content(role="user", parts=parts))
            elif role == "assistant":
                if parts:
                    contents.append(types.Content(role="model", parts=parts))
        
        # 4. Prepare Tools Config
        config_tools = None
        if tools:
            decls = []
            for t in tools:
                if t.get("type") == "function":
                    func = t["function"]
                    decls.append(types.FunctionDeclaration(
                        name=func["name"],
                        description=func.get("description", ""),
                        parameters=func.get("parameters")
                    ))
            if decls:
                config_tools = [types.Tool(function_declarations=decls)]

        # 5. Execute call
        config = types.GenerateContentConfig(
            system_instruction=system_instruction,
            tools=config_tools
        )
        
        try:
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=config
            )
            
            # Parse response parts
            assistant_content = ""
            assistant_tool_calls = []
            
            parts = response.candidates[0].content.parts if response.candidates else []
            
            for part in parts:
                if part.text:
                    assistant_content += part.text
                if part.function_call:
                    fc = part.function_call
                    assistant_tool_calls.append({
                        "id": f"call_{fc.name}",
                        "type": "function",
                        "function": {
                            "name": fc.name,
                            "arguments": json.dumps(fc.args) if fc.args else "{}"
                        }
                    })
            
            if not assistant_content and not assistant_tool_calls:
                 # Check if the SDK returned text via the helper (fallback)
                 # Avoid accessing property unless necessary to prevent warnings
                 sdk_text = getattr(response, 'text', None)
                 if sdk_text:
                     assistant_content = sdk_text
                 else:
                     raise ValueError("Gemini returned an empty response")
            
            res = {
                "role": "assistant",
                "content": assistant_content,
            }
            if assistant_tool_calls:
                res["tool_calls"] = assistant_tool_calls
                
            return res
            
        except Exception as e:
            error_msg = str(e)
            # Only trigger image error if it's actually about media support
            if "image" in error_msg.lower() and ("support" in error_msg.lower() or "type" in error_msg.lower()):
                if "tool" not in error_msg.lower(): # Tool errors are different
                    raise ValueError(f"Gemini model '{self.model_name}' does not support this image/data type.") from e
            
            if "empty" in error_msg.lower() or "must contain" in error_msg.lower():
                raise ValueError(f"Gemini model returned empty response. This may be due to content filtering or model limitations. Original error: {error_msg}")
            raise

    async def chat_stream(
        self, 
        messages: List[Dict[str, Any]], 
        tools: Optional[List[Dict[str, Any]]] = None,
        on_token: Optional[Callable[[str], None]] = None
    ) -> Any:
        from .utils import extract_content
        from google.genai import types
        import json
        import asyncio
        import inspect
        
        # Build contents list (logic same as chat)
        contents = []
        system_instruction = None
        
        for msg in messages:
            role = msg.get("role")
            raw_content = msg.get("content", "")
            tool_calls = msg.get("tool_calls")
            
            if role == "system":
                if isinstance(raw_content, str):
                    system_instruction = raw_content
                elif isinstance(raw_content, list):
                    text_parts = [p.get("text", "") for p in raw_content if p.get("type") == "text"]
                    system_instruction = " ".join(text_parts)
                continue
            
            parts = []
            if raw_content:
                text_content, images = extract_content(raw_content)
                if text_content:
                    parts.append(types.Part.from_text(text=text_content))
                for img in images:
                    if img["data"] and img["mime_type"]:
                        parts.append(types.Part.from_bytes(data=img["data"], mime_type=img["mime_type"]))
            
            if role == "assistant" and tool_calls:
                for tc in tool_calls:
                    func_data = tc.get("function", {})
                    name = func_data.get("name")
                    args_raw = func_data.get("arguments")
                    if isinstance(args_raw, str):
                        try: args = json.loads(args_raw)
                        except Exception: args = {}
                    else: args = args_raw or {}
                    parts.append(types.Part.from_function_call(name=name, args=args))
            
            if role == "tool":
                name = msg.get("name")
                if not name and msg.get("tool_call_id"):
                    tid = msg.get("tool_call_id")
                    if tid.startswith("call_"): name = tid[5:]
                
                tool_content = msg.get("content")
                if isinstance(tool_content, str):
                    try:
                        resp_dict = json.loads(tool_content)
                        if not isinstance(resp_dict, dict): resp_dict = {"result": resp_dict}
                    except Exception: resp_dict = {"result": tool_content}
                else: resp_dict = tool_content or {}
                
                contents.append(types.Content(role="tool", parts=[types.Part.from_function_response(name=name or "unknown_function", response=resp_dict)]))
                continue

            if role == "user":
                if parts: contents.append(types.Content(role="user", parts=parts))
            elif role == "assistant":
                if parts: contents.append(types.Content(role="model", parts=parts))
        
        config_tools = None
        if tools:
            decls = []
            for t in tools:
                if t.get("type") == "function":
                    func = t["function"]
                    decls.append(types.FunctionDeclaration(name=func["name"], description=func.get("description", ""), parameters=func.get("parameters")))
            if decls: config_tools = [types.Tool(function_declarations=decls)]

        config = types.GenerateContentConfig(system_instruction=system_instruction, tools=config_tools)
        
        try:
            # Note: Google GenAI SDK doesn't have an async stream yet in some versions, 
            # but we can wrap it or use the streaming helper.
            # Assuming sync stream for now as per SDK docs usually.
            
            assistant_content = ""
            assistant_tool_calls = []
            
            # Using synchronous stream in thread to avoid blocking event loop
            def get_stream():
                return self.client.models.generate_content_stream(model=self.model_name, contents=contents, config=config)

            loop = asyncio.get_event_loop()
            stream = await loop.run_in_executor(None, get_stream)
            
            for chunk in stream:
                # Extract text
                if chunk.text:
                    token = chunk.text
                    assistant_content += token
                    if on_token:
                        if inspect.iscoroutinefunction(on_token):
                            await on_token(token)
                        else:
                            on_token(token)
                
                # Extract tool calls (usually only in one chunk or first/last)
                if chunk.candidates and chunk.candidates[0].content.parts:
                    for part in chunk.candidates[0].content.parts:
                        if part.function_call:
                            fc = part.function_call
                            assistant_tool_calls.append({
                                "id": f"call_{fc.name}",
                                "type": "function",
                                "function": {
                                    "name": fc.name,
                                    "arguments": json.dumps(fc.args) if fc.args else "{}"
                                }
                            })
            
            res = {"role": "assistant", "content": assistant_content}
            if assistant_tool_calls:
                res["tool_calls"] = assistant_tool_calls
            return res
            
        except Exception as e:
            # Same error handling as chat...
            raise e

    def get_model_name(self) -> str:
        return self.model_name
