"""
Parser for Kimi CLI stream-json output.

Kimi CLI with --print --output-format=stream-json outputs JSONL with chat-completion-like messages:
- {"role":"assistant","content":"Hello!"}
- {"role":"assistant","content":"Let me check.","tool_calls":[{"type":"function","id":"tc_1","function":{"name":"Shell","arguments":"{\"command\":\"ls\"}"}}]}
- {"role":"tool","tool_call_id":"tc_1","content":"result"}
"""

import time
import json
from typing import Any, Dict, Optional, List, 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


# Kimi tool name mapping (for notifications)
KIMI_TOOL_NAMES = {
    "Shell": "Shell",
    "Read": "Read",
    "Write": "Write",
    "Edit": "Edit",
    "Grep": "Grep",
    "Glob": "Glob",
    "TodoWrite": "TodoWrite",
    "TodoRead": "TodoRead",
    "Bash": "Bash",
    "ListFiles": "ListFiles",
    "SearchFiles": "SearchFiles",
}


def is_assistant_message(line_dict: Dict[str, Any]) -> bool:
    """Check if the JSONL line is an assistant message."""
    return line_dict.get("role") == "assistant"


def is_tool_result(line_dict: Dict[str, Any]) -> bool:
    """Check if the JSONL line is a tool result."""
    return line_dict.get("role") == "tool"


def normalize_content(content: Any) -> Optional[str]:
    """
    Normalize Kimi content to a string.

    Kimi CLI can return content in two formats:
    1. Simple string: "Hello!"
    2. List of content blocks: [{"type": "think", "think": "..."}, {"type": "text", "text": "..."}]

    This function extracts text content from both formats.
    """
    if content is None:
        return None

    if isinstance(content, str):
        return content

    if isinstance(content, list):
        # Extract text from content blocks
        text_parts = []
        for block in content:
            if isinstance(block, dict):
                # Handle text blocks
                if block.get("type") == "text":
                    text_parts.append(block.get("text", ""))
                # Note: We skip "think" blocks as they are internal reasoning
        return "\n".join(text_parts) if text_parts else None

    return None


def get_text_content(line_dict: Dict[str, Any]) -> Optional[str]:
    """Extract text content from an assistant message."""
    if not is_assistant_message(line_dict):
        return None
    return normalize_content(line_dict.get("content"))


def extract_tool_info(line_dict: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]:
    """
    Extract tool call information from an assistant message.

    Returns:
        List of dicts with tool_name, arguments, id; or None if no tool calls.
    """
    if not is_assistant_message(line_dict):
        return None

    tool_calls = line_dict.get("tool_calls")
    if not tool_calls:
        return None

    result = []
    for tc in tool_calls:
        if tc.get("type") != "function":
            continue
        func = tc.get("function", {})
        tool_name = func.get("name", "unknown")
        arguments_str = func.get("arguments", "{}")
        try:
            arguments = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
        except json.JSONDecodeError:
            arguments = {"raw": arguments_str}

        result.append({
            "tool_name": tool_name,
            "arguments": arguments,
            "id": tc.get("id", ""),
        })

    return result if result else None


def to_model_response(line_dict: Dict[str, Any], *, default_model: str = "kimi-model") -> "ModelResponse":
    """
    Convert a Kimi CLI JSONL line to a LiteLLM ModelResponse.

    Handles:
    - Assistant messages with text content
    - Assistant messages with tool_calls
    - Tool result messages (role=tool)
    """
    from litellm.utils import ModelResponse
    from litellm.types.utils import Message, ChatCompletionMessageToolCall, Function

    role = line_dict.get("role", "assistant")
    # Normalize content to handle both string and list formats
    content = normalize_content(line_dict.get("content"))

    # Kimi CLI doesn't provide explicit final markers; rely on on_complete
    has_tool_calls = bool(line_dict.get("tool_calls"))
    is_final = False

    blocks_provider_specific_fields = {
        BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE: is_final
    }

    # Build tool calls if present
    tool_calls_list = None
    if has_tool_calls and is_assistant_message(line_dict):
        tool_calls_list = []
        for tc in line_dict.get("tool_calls", []):
            if tc.get("type") != "function":
                continue
            func = tc.get("function", {})
            tool_calls_list.append(
                ChatCompletionMessageToolCall(
                    id=tc.get("id", f"call_{int(time.time())}"),
                    type="function",
                    function=Function(
                        name=func.get("name", "unknown"),
                        arguments=func.get("arguments", "{}"),
                    ),
                )
            )

    # Build ModelResponse
    mr = ModelResponse()
    mr.choices[0].message = Message(
        role="assistant" if role == "assistant" else "assistant",
        content=content,
        tool_calls=tool_calls_list if tool_calls_list else None,
        provider_specific_fields=blocks_provider_specific_fields,
    )
    mr.choices[0].finish_reason = "tool_calls" if tool_calls_list else "stop"
    mr.model = default_model
    mr.created = int(time.time())
    mr.usage = {}

    return mr
