import time
import json
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from blocks_control_sdk.parsers.base_messages import BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE

if TYPE_CHECKING:
    from litellm.utils import ModelResponse


# https://github.com/BerriAI/litellm/blob/main/litellm/llms/gemini/chat/transformation.py#L70

GEMINI_TOOLS__READ_FILE = "read_file" # Args: {"absolute_path": "/bundle/workspace/bundle/tools/register.py"}
GEMINI_TOOLS__READ_MANY_FILES = "read_many_files" # Args: {"paths": ["/bundle/workspace/bundle/tools/register.py", "/bundle/workspace/bundle/tools/blocks.py"]}
GEMINI_TOOLS__LIST_DIRECTORY = "list_directory" # Args: {"path": "/bundle/workspace/bundle/"}
GEMINI_TOOLS__RUN_SHELL_COMMAND = "run_shell_command" # Args: {"description": "...", "command": "..."}
GEMINI_TOOLS__WRITE_FILE = "write_file" # Args: {"content": "...", "file_path": "..."}


def _extract_text_from_parts(parts: List[Dict[str, Any]]) -> Optional[str]:
    """Extract text from Gemini parts array format."""
    if not parts or not isinstance(parts, list):
        return None
    
    texts = []
    for part in parts:
        if isinstance(part, dict) and "text" in part:
            text = part.get("text", "")
            if text:
                texts.append(text)
    
    return "\n".join(texts) if texts else None


def _extract_function_calls_from_parts(parts: List[Dict[str, Any]]) -> Optional[List[Dict[str, Any]]]:
    """Extract function calls from Gemini parts array format."""
    if not parts or not isinstance(parts, list):
        return None
    
    function_calls = []
    for part in parts:
        if isinstance(part, dict) and "functionCall" in part:
            func_call = part["functionCall"]
            # Convert Gemini function call format to OpenAI format
            function_calls.append({
                "id": f"call_{int(time.time() * 1000)}",
                "type": "function",
                "function": {
                    "name": func_call.get("name", ""),
                    "arguments": json.dumps(func_call.get("args", {})) if isinstance(func_call.get("args"), dict) else "{}"
                }
            })
    
    return function_calls if function_calls else None


def _extract_thinking_from_parts(parts: List[Dict[str, Any]]) -> Optional[str]:
    """Extract thinking/reasoning content from Gemini parts array format."""
    if not parts or not isinstance(parts, list):
        return None
    
    thinking_texts = []
    for part in parts:
        if isinstance(part, dict) and "thought" in part:
            thought = part.get("thought", "")
            if thought:
                thinking_texts.append(thought)
        elif isinstance(part, dict) and "thoughtSignature" in part:
            # Handle thoughtSignature as a thinking block marker
            thinking_texts.append("[Thinking...]")
    
    return "\n".join(thinking_texts) if thinking_texts else None


def _normalize_blob_to_openai_format(blob: Dict[str, Any], *, default_model: str = "unknown-model") -> Dict[str, Any]:
    """Normalize different Gemini blob types to standard OpenAI API response format."""
    
    blob_type = blob.get("type")
    blob_id = f"msg_{int(time.time())}"
    role = blob.get("role")
    
    # Handle native Gemini format: {"role": "model", "parts": [...]}
    if role and "parts" in blob:
        parts = blob.get("parts", [])
        
        # Extract content from parts
        text_content = _extract_text_from_parts(parts)
        function_calls = _extract_function_calls_from_parts(parts)
        thinking_content = _extract_thinking_from_parts(parts)
        
        # Map Gemini role to OpenAI role
        openai_role = "assistant" if role == "model" else role
        
        message = {
            "role": openai_role,
            "content": text_content
        }
        
        if function_calls:
            message["tool_calls"] = function_calls
        
        if thinking_content:
            message["reasoning_content"] = thinking_content
        
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": message,
                "finish_reason": "tool_calls" if function_calls else "stop"
            }],
            "usage": blob.get("usage", {})
        }
    
    # Handle wrapped result format
    elif blob_type == "result":
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": blob.get("result", "")
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }

    # Handle final response format
    elif blob_type == "final":
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": blob.get("response", "")
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }

    elif blob_type == "message" and role == "assistant":
        content_text = blob.get("content", "")
        delta = blob.get("delta", {})
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": content_text
                },
                "provider_specific_fields": {
                    "delta": delta
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }
    # Handle content response format
    elif blob_type == "content":
        content_text = blob.get("content", "")
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": content_text
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }

    # Handle assistant wrapper format
    elif blob_type == "assistant" and "message" in blob:
        message = blob["message"]
        # If the message has Gemini structure, process it
        if "parts" in message:
            parts = message.get("parts", [])
            text_content = _extract_text_from_parts(parts)
            function_calls = _extract_function_calls_from_parts(parts)
            thinking_content = _extract_thinking_from_parts(parts)
            
            processed_message = {
                "role": "assistant",
                "content": text_content
            }
            
            if function_calls:
                processed_message["tool_calls"] = function_calls
            
            if thinking_content:
                processed_message["reasoning_content"] = thinking_content
            
            return {
                "id": blob_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": message.get("model", default_model),
                "choices": [{
                    "index": 0,
                    "message": processed_message,
                    "finish_reason": "tool_calls" if function_calls else "stop"
                }],
                "usage": message.get("usage", {})
            }
        else:
            # Already in OpenAI format
            return {
                "id": blob_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": message.get("model", default_model),
                "choices": [{
                    "index": 0,
                    "message": message,
                    "finish_reason": "stop"
                }],
                "usage": message.get("usage", {})
            }
    
    # Handle Gemini streaming format with candidates
    elif "candidates" in blob:
        candidates = blob.get("candidates", [])
        if candidates and len(candidates) > 0:
            candidate = candidates[0]
            content = candidate.get("content", {})
            parts = content.get("parts", [])
            
            text_content = _extract_text_from_parts(parts)
            function_calls = _extract_function_calls_from_parts(parts)
            thinking_content = _extract_thinking_from_parts(parts)
            
            message = {
                "role": content.get("role", "assistant"),
                "content": text_content
            }
            
            if function_calls:
                message["tool_calls"] = function_calls
            
            if thinking_content:
                message["reasoning_content"] = thinking_content
            
            finish_reason = candidate.get("finishReason", "STOP")
            # Map Gemini finish reasons to OpenAI
            finish_reason_map = {
                "STOP": "stop",
                "MAX_TOKENS": "length",
                "SAFETY": "content_filter",
                "RECITATION": "content_filter",
                "OTHER": "stop"
            }
            
            return {
                "id": blob_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": blob.get("model", default_model),
                "choices": [{
                    "index": candidate.get("index", 0),
                    "message": message,
                    "finish_reason": finish_reason_map.get(finish_reason, "stop")
                }],
                "usage": blob.get("usageMetadata", {})
            }

    # Handle telemetry events
    elif blob_type == "tool_use":
        function_name = blob.get("tool_name", "")
        function_args = blob.get("parameters", {})
        function_id = blob.get("tool_id", "")
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [{
                        "id": function_id,
                        "type": "function",
                        "function": {
                            "name": function_name,
                            "arguments": json.dumps(function_args) if isinstance(function_args, dict) else "{}"
                        }
                    }]
                },
                "finish_reason": "tool_calls"
            }],
            "usage": blob.get("usage", {})
        }
    elif blob_type == "telemetry":
        event = blob.get("event", {})
        event_name = event.get("event.name", "")

        # Handle tool call telemetry events
        if event_name == "tool_call":
            function_name = event.get("function_name", "")
            function_args = event.get("function_args", {})

            if function_name:
                # Create OpenAI-style tool call
                tool_call = {
                    "id": f"call_{int(time.time() * 1000)}",
                    "type": "function",
                    "function": {
                        "name": function_name,
                        "arguments": json.dumps(function_args) if isinstance(function_args, dict) else "{}"
                    }
                }

                message = {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [tool_call]
                }

                return {
                    "id": blob_id,
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": blob.get("model", default_model),
                    "choices": [{
                        "index": 0,
                        "message": message,
                        "finish_reason": "tool_calls"
                    }],
                    "usage": blob.get("usage", {})
                }
            else:
                # Tool call event without function name - return empty message
                return {
                    "id": blob_id,
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": blob.get("model", default_model),
                    "choices": [{
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": ""
                        },
                        "finish_reason": "stop"
                    }],
                    "usage": blob.get("usage", {})
                }
        else:
            # Non-tool-call telemetry event - return empty message
            return {
                "id": blob_id,
                "object": "chat.completion",
                "created": int(time.time()),
                "model": blob.get("model", default_model),
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": ""
                    },
                    "finish_reason": "stop"
                }],
                "usage": blob.get("usage", {})
            }

    else:
        raise ValueError(f"Unsupported blob format: {json.dumps(blob, indent=2)}")


def to_model_response(blob: Dict[str, Any], *, default_model: str = "unknown-model") -> "ModelResponse":
    from litellm.utils import ModelResponse, convert_to_model_response_object

    """
    Convert various Gemini blob formats to LiteLLM ModelResponse.
    
    Supports:
    - Native Gemini format: {"role": "model", "parts": [...]}
    - Gemini streaming format: {"candidates": [...]}
    - type: "result" - Legacy result format
    - type: "assistant" - Legacy assistant format
    """
    
    # Check if this is already a standard OpenAI response (no 'type' field but has 'choices')
    if "type" not in blob and "choices" in blob:
        # This is already in standard OpenAI format, use it directly
        normalized_response = blob
    else:
        # Normalize Gemini formats to standard OpenAI format
        normalized_response = _normalize_blob_to_openai_format(blob, default_model=default_model)
    
    # Create a base ModelResponse object
    model_response = ModelResponse()
    model_response.model = normalized_response.get("model", default_model)
    
    # Set provider-specific fields for blocks
    is_final_message = blob.get("type") in ["result", "final"]  # Both result and final types are considered final
    hidden_params = {
        "original_response": blob,
        BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE: is_final_message
    }
    
    # Use litellm's utility to do the heavy lifting
    final_response = convert_to_model_response_object(
        response_object=normalized_response,
        model_response_object=model_response,
        response_type="completion",
        hidden_params=hidden_params
    )
    
    # Ensure provider-specific fields are set on the message
    if final_response.choices and final_response.choices[0].message:
        if not final_response.choices[0].message.provider_specific_fields:
            final_response.choices[0].message.provider_specific_fields = {}
        final_response.choices[0].message.provider_specific_fields[BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE] = is_final_message
    
    return final_response

