from typing import Any

from pydantic import Field

from dhenara.ai.types.shared.base import BaseModel
from dhenara.ai.utils.dai_disk import DAI_JSON


class ChatResponseToolCall(BaseModel):
    """Representation of a tool call from an LLM"""

    call_id: str | None = Field(
        default=None,
        description="An identifier used to map this tool call to a tool call output.",
    )
    id: str | None = Field(
        default=None,
        description="The unique ID of the tool call output in the provider platform.",
    )
    name: str
    arguments: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    raw_data: str | dict | None = Field(
        default=None,
        description="Raw unparsed response from the model",
    )
    parse_error: str | None = Field(
        default=None,
        description="Error that occurred during parsing, if any",
    )

    @classmethod
    def parse_args_str_or_dict(cls, arguments: str | dict) -> dict:
        """Parse arguments from either JSON striing or dict"""
        arguments_dict = {}
        raw_data = None
        parse_error = None
        if isinstance(arguments, str):
            try:
                arguments_dict = DAI_JSON.loads(arguments)
            except Exception as e:
                raw_data = arguments
                parse_error = str(e)
        else:
            arguments_dict = arguments

        return {
            "arguments_dict": arguments_dict,
            "raw_data": raw_data,
            "parse_error": parse_error,
        }


class ChatResponseToolCallResult(BaseModel):
    """Result of executing a tool call, which may pass to LLM in next turn"""

    tool_name: str
    call_id: str | None = None
    result: Any = None
    error: str | None = None
