import json
import logging
import warnings
from dataclasses import dataclass, field, is_dataclass
from typing import Any, Dict, List, Literal, Optional, Union
from uuid import UUID, uuid4

from requests import HTTPError

from freeplay import api_support
from freeplay.errors import FreeplayClientError, FreeplayError
from freeplay.llm_parameters import LLMParameters
from freeplay.model import (
    InputVariables,
    MediaInputMap,
    NormalizedOutputSchema,
    OpenAIFunctionCall,
    TestRunInfo,
)
from freeplay.resources.prompts import (
    PromptInfo,
    PromptVersionInfo,
)
from freeplay.resources.sessions import SessionInfo, TraceInfo
from freeplay.support import CallSupport, media_inputs_to_json
from freeplay.utils import convert_provider_message_to_dict

logger = logging.getLogger(__name__)


@dataclass
class UsageTokens:
    prompt_tokens: int
    completion_tokens: int


ApiStyle = Union[Literal["batch"], Literal["default"]]


@dataclass
class CallInfo:
    provider: Optional[str] = None
    model: Optional[str] = None
    start_time: Optional[float] = None
    end_time: Optional[float] = None
    model_parameters: Optional[LLMParameters] = None
    provider_info: Optional[Dict[str, Any]] = None
    usage: Optional[UsageTokens] = None
    api_style: Optional[ApiStyle] = None

    @staticmethod
    def from_prompt_info(
        prompt_info: PromptInfo,
        start_time: float,
        end_time: float,
        usage: Optional[UsageTokens] = None,
        api_style: Optional[ApiStyle] = None,
    ) -> "CallInfo":
        return CallInfo(
            provider=prompt_info.provider,
            model=prompt_info.model,
            start_time=start_time,
            end_time=end_time,
            model_parameters=prompt_info.model_parameters,
            provider_info=prompt_info.provider_info,
            usage=usage,
            api_style=api_style,
        )


@dataclass
class ResponseInfo:
    is_complete: Optional[bool] = None
    function_call_response: Optional[OpenAIFunctionCall] = None
    prompt_tokens: Optional[int] = None
    response_tokens: Optional[int] = None


@dataclass
class RecordPayload:
    project_id: str
    all_messages: List[Dict[str, Any]]

    session_info: SessionInfo = field(
        default_factory=lambda: SessionInfo(
            session_id=str(uuid4()), custom_metadata=None
        )
    )
    inputs: Optional[InputVariables] = None
    prompt_version_info: Optional[PromptVersionInfo] = None
    call_info: Optional[CallInfo] = None
    media_inputs: Optional[MediaInputMap] = None
    # Accepts dicts in provider-specific tool schema format (OpenAI, Anthropic, GenAI/Vertex)
    # Backend automatically normalizes all formats
    tool_schema: Optional[List[Dict[str, Any]]] = None
    output_schema: Optional[NormalizedOutputSchema] = None
    response_info: Optional[ResponseInfo] = None
    test_run_info: Optional[TestRunInfo] = None
    eval_results: Optional[Dict[str, Union[bool, float]]] = None
    parent_id: Optional[UUID] = None
    completion_id: Optional[UUID] = None
    # Deprecated field support for backward compatibility
    # This field will be removed in v0.6.0
    trace_info: Optional[TraceInfo] = None


@dataclass
class RecordUpdatePayload:
    project_id: str
    completion_id: str
    new_messages: Optional[List[Dict[str, Any]]] = None
    eval_results: Optional[Dict[str, Union[bool, float]]] = None


@dataclass
class RecordResponse:
    completion_id: str


class Recordings:
    def __init__(self, call_support: CallSupport):
        self.call_support = call_support

    @staticmethod
    def _contains_bytes(obj: Any) -> bool:
        """Check if an object contains bytes that need conversion."""
        if isinstance(obj, bytes):
            return True
        elif isinstance(obj, dict):
            return any(Recordings._contains_bytes(v) for v in obj.values())  # pyright: ignore[reportUnknownVariableType]
        elif isinstance(obj, list):
            return any(Recordings._contains_bytes(item) for item in obj)  # pyright: ignore[reportUnknownVariableType]
        elif is_dataclass(obj):
            # Check dataclass fields for bytes
            for field_name in obj.__dataclass_fields__:
                if Recordings._contains_bytes(getattr(obj, field_name)):
                    return True
        return False

    def create(self, record_payload: RecordPayload) -> RecordResponse:
        if len(record_payload.all_messages) < 1:
            raise FreeplayClientError(
                "Messages list must have at least one message. "
                "The last message should be the current response."
            )

        if record_payload.tool_schema is not None:
            record_payload.tool_schema = [
                convert_provider_message_to_dict(tool)
                for tool in record_payload.tool_schema
            ]

        # Convert messages if using Bedrock provider or if messages contain bytes
        messages = record_payload.all_messages
        needs_conversion = (
            record_payload.call_info and record_payload.call_info.provider == "bedrock"
        ) or any(self._contains_bytes(msg) for msg in messages)
        if needs_conversion:
            messages = [convert_provider_message_to_dict(msg) for msg in messages]

        record_api_payload: Dict[str, Any] = {
            "messages": messages,
            "inputs": record_payload.inputs,
            "tool_schema": record_payload.tool_schema,
            "output_schema": record_payload.output_schema,
            "session_info": {
                "custom_metadata": record_payload.session_info.custom_metadata
            },
            "parent_id": str(record_payload.parent_id)
            if record_payload.parent_id is not None
            else None,
        }

        if record_payload.prompt_version_info is not None:
            record_api_payload["prompt_info"] = {
                "environment": record_payload.prompt_version_info.environment,
                "prompt_template_version_id": record_payload.prompt_version_info.prompt_template_version_id,
            }

        if record_payload.call_info is not None:
            record_api_payload["call_info"] = {
                "start_time": record_payload.call_info.start_time,
                "end_time": record_payload.call_info.end_time,
                "model": record_payload.call_info.model,
                "provider": record_payload.call_info.provider,
                "provider_info": record_payload.call_info.provider_info,
                "llm_parameters": record_payload.call_info.model_parameters,
                "api_style": record_payload.call_info.api_style,
            }

        if record_payload.completion_id is not None:
            record_api_payload["completion_id"] = str(record_payload.completion_id)

        if record_payload.session_info.custom_metadata is not None:
            record_api_payload["custom_metadata"] = (
                record_payload.session_info.custom_metadata
            )

        if record_payload.response_info is not None:
            if record_payload.response_info.function_call_response is not None:
                record_api_payload["response_info"] = {
                    "function_call_response": {
                        "name": record_payload.response_info.function_call_response[
                            "name"
                        ],
                        "arguments": record_payload.response_info.function_call_response[
                            "arguments"
                        ],
                    }
                }

        if record_payload.test_run_info is not None:
            record_api_payload["test_run_info"] = {
                "test_run_id": record_payload.test_run_info.test_run_id,
                "test_case_id": record_payload.test_run_info.test_case_id,
            }

        if record_payload.eval_results is not None:
            record_api_payload["eval_results"] = record_payload.eval_results

        if record_payload.trace_info is not None:
            warnings.warn(
                "trace_info in RecordPayload is deprecated and will be removed in v0.6.0. Use parent_id instead.",
                DeprecationWarning,
            )
            record_api_payload["trace_info"] = {
                "trace_id": record_payload.trace_info.trace_id
            }

        if (
            record_payload.call_info is not None
            and record_payload.call_info.usage is not None
        ):
            record_api_payload["call_info"]["usage"] = {
                "prompt_tokens": record_payload.call_info.usage.prompt_tokens,
                "completion_tokens": record_payload.call_info.usage.completion_tokens,
            }

        if record_payload.media_inputs is not None:
            record_api_payload["media_inputs"] = {
                name: media_inputs_to_json(media_input)
                for name, media_input in record_payload.media_inputs.items()
            }

        try:
            recorded_response = api_support.post_raw(
                api_key=self.call_support.freeplay_api_key,
                url=f"{self.call_support.api_base}/v2/projects/{record_payload.project_id}/sessions/{record_payload.session_info.session_id}/completions",
                payload=record_api_payload,
            )
            recorded_response.raise_for_status()
            json_dom = recorded_response.json()
            return RecordResponse(completion_id=str(json_dom["completion_id"]))
        except HTTPError as e:
            message = (
                f"There was an error recording to Freeplay. Call will not be logged. "
                f"Status: {e.response.status_code}. "
            )

            self.__handle_and_raise_api_error(e, message)

        except Exception as e:
            status_code = -1
            if hasattr(e, "response"):
                response = getattr(e, "response")
                if hasattr(response, "status_code"):
                    status_code = getattr(response, "status_code")

            message = (
                f"There was an error recording to Freeplay. Call will not be logged. "
                f"Status: {status_code}. {e.__class__}"
            )

            raise FreeplayError(message) from e

        raise FreeplayError("Unexpected error occurred while recording to Freeplay.")

    def update(self, record_update_payload: RecordUpdatePayload) -> RecordResponse:  # type: ignore
        # Only convert messages if they contain bytes (we don't have provider info for updates)
        new_messages = record_update_payload.new_messages
        if new_messages and any(self._contains_bytes(msg) for msg in new_messages):
            new_messages = [
                convert_provider_message_to_dict(msg) for msg in new_messages
            ]

        record_update_api_payload: Dict[str, Any] = {
            "new_messages": new_messages,
            "eval_results": record_update_payload.eval_results,
        }

        try:
            record_update_response = api_support.post_raw(
                api_key=self.call_support.freeplay_api_key,
                url=f"{self.call_support.api_base}/v2/projects/{record_update_payload.project_id}/completions/{record_update_payload.completion_id}",
                payload=record_update_api_payload,
            )
            record_update_response.raise_for_status()
            json_dom = record_update_response.json()
            return RecordResponse(completion_id=str(json_dom["completion_id"]))
        except HTTPError as e:
            message = f"There was an error updating the completion. Status: {e.response.status_code}."
            self.__handle_and_raise_api_error(e, message)

    @staticmethod
    def __handle_and_raise_api_error(e: HTTPError, messages: str) -> None:
        if e.response.content:
            try:
                content = e.response.content
                json_body = json.loads(content)
                if "message" in json_body:
                    messages += json_body["message"]
            except Exception:
                pass
        else:
            messages += f"{e.__class__}"
        raise FreeplayError(messages) from e
