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


OPENAI_TOOLS__SHELL = "shell" # Args: {"command": ["bash", "-c", "..."]}


def _extract_text_from_content_array(content_array: List[Dict[str, Any]]) -> Optional[str]:
    """Extract text from OpenAI content array format."""
    if not content_array or not isinstance(content_array, list):
        return None
    
    texts = []
    for item in content_array:
        if isinstance(item, dict) and item.get("type") == "output_text":
            text = item.get("text", "")
            if text:
                texts.append(text)
    
    return "\n".join(texts) if texts else None


def _normalize_blob_to_openai_format(blob: Dict[str, Any], *, default_model: str = "unknown-model") -> Dict[str, Any]:
    """Normalize different blob types to standard OpenAI API response format."""
    
    blob_type = blob.get("type")
    blob_id = blob.get("id", f"msg_{int(time.time())}")
    
    if blob_type == "message":
        # Handle: {"type":"message","id":"msg_...","role":"assistant","content":[{"type":"output_text","text":"..."}]}
        role = blob.get("role")
        if role != "assistant":
            raise ValueError(f"Expected message role 'assistant', got '{role}'")
        
        content = blob.get("content")
        if isinstance(content, list):
            # Extract text from content array
            text_content = _extract_text_from_content_array(content)
        else:
            text_content = 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": text_content
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }
    
    elif blob_type == "function_call":
        # Handle: {"type":"function_call","id":"fc_...","name":"...","arguments":"...","call_id":"..."}
        function_name = blob.get("name")
        function_args = blob.get("arguments", "{}")
        call_id = blob.get("call_id", blob_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": call_id,
                        "type": "function",
                        "function": {
                            "name": function_name,
                            "arguments": function_args
                        }
                    }]
                },
                "finish_reason": "tool_calls"
            }],
            "usage": blob.get("usage", {})
        }
    elif blob_type == "function_call_output":
        # Handle tool output
        tool_call_id = blob.get("tool_call_id") or blob.get("call_id")
        content = blob.get("content") or blob.get("output")
        name = blob.get("name", "uknown")
        
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "tool",
                    "content": content,
                    "tool_call_id": tool_call_id,
                    "name": name
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }
    elif blob_type == "reasoning":
        # Handle: {"type":"reasoning","id":"rs_...","summary":[],"content":null,"encrypted_content":null}
        reasoning_content = blob.get("content") or blob.get("encrypted_content")
        summary = blob.get("summary", [])
        
        # Convert summary array to text if needed
        if isinstance(summary, list) and summary:
            # check if the items are dicts or strings
            if all(isinstance(item, dict) for item in summary):
                summary_text = ".\n ".join(item.get("text", "") if item.get("summary_text") == "summary_text" else "" for item in summary)
            else:
                summary_text = "\n".join(str(item) for item in summary)
        else:
            summary_text = "Reasoning..."
        
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": summary_text,
                    "reasoning_content": reasoning_content
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }
    
    elif blob_type == "result":
        # Handle legacy format: {"type": "result", "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", {})
        }
    
    elif blob_type == "assistant" and "message" in blob:
        # Handle legacy format: {"type":"assistant","message":{...}}
        message = blob["message"]
        role = message.get("role")
        if role != "assistant":
            raise ValueError(f"Expected message role 'assistant', got '{role}'")
        
        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", {})
        }

    elif blob_type == "agent_reasoning":
        message = blob.get("text")
        return {
            "id": blob_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": blob.get("model", default_model),
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": message
                },
                "finish_reason": "stop"
            }],
            "usage": blob.get("usage", {})
        }
    
    else:
        raise ValueError(f"Unsupported blob type: {blob_type}")


def to_model_response(blob: Dict[str, Any], *, default_model: str = "unknown-model") -> "ModelResponse":
    """
    Convert various blob formats to LiteLLM ModelResponse using litellm's convert_to_model_response_object utility.
    
    Supports:
    - Standard OpenAI API responses (no type field)
    - type: "message" - Codex/O3 message format
    - type: "function_call" - Function/tool call format  
    - type: "reasoning" - Reasoning content format
    - type: "result" - Legacy result format
    - type: "assistant" - Legacy assistant format
    """
    from litellm.utils import ModelResponse, convert_to_model_response_object

    # {"timestamp":"2025-09-22T20:10:48.754Z","type":"response_item","payload":{"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hey! I’m here and ready to help—what do you need?"}]}}

    if blob.get("type") == "response_item" or blob.get("type") == "event_msg":
        blob = blob.get("payload")

    # 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 custom blob 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") == "result"  # Only result type is 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
