"""
Supports writing files to Google AI Studio Files API.

For vertex ai, check out the vertex_ai/files/handler.py file.
"""

import time
from typing import Any, List, Literal, Optional
from urllib.parse import urlparse

import httpx
from openai.types.file_deleted import FileDeleted

from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
from litellm.llms.base_llm.files.transformation import (
    BaseFilesConfig,
    LiteLLMLoggingObj,
)
from litellm.types.llms.gemini import GeminiCreateFilesResponseObject
from litellm.types.llms.openai import (
    AllMessageValues,
    CreateFileRequest,
    HttpxBinaryResponseContent,
    OpenAICreateFileRequestOptionalParams,
    OpenAIFileObject,
)
from litellm.types.utils import LlmProviders

from ..common_utils import GeminiModelInfo


class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
    def __init__(self):
        pass

    @property
    def custom_llm_provider(self) -> LlmProviders:
        return LlmProviders.GEMINI

    def validate_environment(
        self,
        headers: dict[Any, Any],
        model: str,
        messages: List[AllMessageValues],
        optional_params: dict[Any, Any],
        litellm_params: dict[Any, Any],
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
    ) -> dict[Any, Any]:
        """
        Validate environment and add Gemini API key to headers.
        Google AI Studio uses x-goog-api-key header for authentication.
        """
        resolved_api_key = self.get_api_key(api_key)
        if not resolved_api_key:
            raise ValueError(
                "GEMINI_API_KEY is required for Google AI Studio file operations"
            )

        headers["x-goog-api-key"] = resolved_api_key
        return headers

    def get_complete_url(
        self,
        api_base: Optional[str],
        api_key: Optional[str],
        model: str,
        optional_params: dict,
        litellm_params: dict,
        stream: Optional[bool] = None,
    ) -> str:
        """
        OPTIONAL

        Get the complete url for the request

        Some providers need `model` in `api_base`
        """
        endpoint = "upload/v1beta/files"
        api_base = self.get_api_base(api_base)
        if not api_base:
            raise ValueError("api_base is required")

        # Get API key from multiple sources
        final_api_key = api_key or litellm_params.get("api_key") or self.get_api_key()
        if not final_api_key:
            raise ValueError("api_key is required")

        url = "{}/{}".format(api_base, endpoint)
        return url

    def get_supported_openai_params(
        self, model: str
    ) -> List[OpenAICreateFileRequestOptionalParams]:
        return []

    def map_openai_params(
        self,
        non_default_params: dict,
        optional_params: dict,
        model: str,
        drop_params: bool,
    ) -> dict:
        return optional_params

    def transform_create_file_request(
        self,
        model: str,
        create_file_data: CreateFileRequest,
        optional_params: dict,
        litellm_params: dict,
    ) -> dict:
        """
        Transform the OpenAI-style file creation request into Gemini's format

        Returns:
            dict: Contains both request data and headers for the two-step upload
        """
        # Extract the file information
        file_data = create_file_data.get("file")
        if file_data is None:
            raise ValueError("File data is required")

        # Use the common utility function to extract file data
        extracted_data = extract_file_data(file_data)

        # Get file size
        file_size = len(extracted_data["content"])

        # Step 1: Initial resumable upload request
        headers = {
            "X-Goog-Upload-Protocol": "resumable",
            "X-Goog-Upload-Command": "start",
            "X-Goog-Upload-Header-Content-Length": str(file_size),
            "X-Goog-Upload-Header-Content-Type": extracted_data["content_type"],
            "Content-Type": "application/json",
        }
        headers.update(extracted_data["headers"])  # Add any custom headers

        # Initial metadata request body
        initial_data = {
            "file": {
                "display_name": extracted_data["filename"] or str(int(time.time()))
            }
        }

        # Step 2: Actual file upload data
        upload_headers = {
            "Content-Length": str(file_size),
            "X-Goog-Upload-Offset": "0",
            "X-Goog-Upload-Command": "upload, finalize",
        }

        return {
            "initial_request": {"headers": headers, "data": initial_data},
            "upload_request": {
                "headers": upload_headers,
                "data": extracted_data["content"],
            },
        }

    def transform_create_file_response(
        self,
        model: Optional[str],
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
    ) -> OpenAIFileObject:
        """
        Transform Gemini's file upload response into OpenAI-style FileObject
        """
        try:
            response_json = raw_response.json()

            response_object = GeminiCreateFilesResponseObject(
                **response_json.get("file", {})  # type: ignore
            )

            # Extract file information from Gemini response

            return OpenAIFileObject(
                id=response_object["uri"],  # Gemini uses URI as identifier
                bytes=int(
                    response_object["sizeBytes"]
                ),  # Gemini doesn't return file size
                created_at=int(
                    time.mktime(
                        time.strptime(
                            response_object["createTime"].replace("Z", "+00:00"),
                            "%Y-%m-%dT%H:%M:%S.%f%z",
                        )
                    )
                ),
                filename=response_object["displayName"],
                object="file",
                purpose="user_data",  # Default to assistants as that's the main use case
                status="uploaded",
                status_details=None,
            )
        except Exception as e:
            verbose_logger.exception(f"Error parsing file upload response: {str(e)}")
            raise ValueError(f"Error parsing file upload response: {str(e)}")

    def transform_retrieve_file_request(
        self,
        file_id: str,
        optional_params: dict,
        litellm_params: dict,
    ) -> tuple[str, dict]:
        """
        Get the URL to retrieve a file from Google AI Studio.

        Endpoint:
        GET https://generativelanguage.googleapis.com/v1beta/{name=files/*}

        The URL should look like:
        https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY

        We expect file_id to be just the file identifier (e.g., files/abc123 or abc123)
        as returned by the upload response. (If it's a full URL, extract the file name.)
        """
        api_key = litellm_params.get("api_key") or self.get_api_key()
        if not api_key:
            raise ValueError("api_key is required")

        file_part = self._normalize_gemini_file_id(file_id)

        api_base = (
            self.get_api_base(litellm_params.get("api_base"))
            or "https://generativelanguage.googleapis.com"
        )
        api_base = api_base.rstrip("/")

        url = f"{api_base}/v1beta/{file_part}"

        # API key is passed via x-goog-api-key header (set in validate_environment)
        return url, {}

    def _normalize_gemini_file_id(self, file_id: str) -> str:
        """
        Normalize file identifier into `files/{id}` form.

        Supports:
        - `abc123`
        - `files/abc123`
        - `https://generativelanguage.googleapis.com/v1beta/files/abc123`
        """
        if file_id.startswith(("http://", "https://")):
            parsed = urlparse(file_id)
            path = parsed.path.lstrip("/")
            files_index = path.find("files/")
            if files_index != -1:
                normalized_file_id = path[files_index:]
            else:
                normalized_file_id = path
        else:
            normalized_file_id = file_id

        normalized_file_id = normalized_file_id.strip("/")
        if normalized_file_id.startswith("files/"):
            normalized_file_id = normalized_file_id.removeprefix("files/")

        encoded_file_id = encode_url_path_segment(
            normalized_file_id, field_name="file_id"
        )

        return f"files/{encoded_file_id}"

    def transform_retrieve_file_response(
        self,
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
    ) -> OpenAIFileObject:
        """
        Transform Gemini's file retrieval response into OpenAI-style FileObject
        """
        try:
            verbose_logger.debug(f"Retrieve file response: {raw_response.text}")
            response_json = raw_response.json()
            verbose_logger.debug(f"Response JSON: {response_json}")
            # Map Gemini state to OpenAI status
            gemini_state = response_json.get("state", "STATE_UNSPECIFIED")
            # Explicitly type status as the Literal union
            if gemini_state == "ACTIVE":
                status: Literal["uploaded", "processed", "error"] = "processed"
            elif gemini_state == "FAILED":
                status = "error"
            else:
                status = "uploaded"

            return OpenAIFileObject(
                id=response_json.get("uri", ""),
                bytes=int(response_json.get("sizeBytes", 0)),
                created_at=int(
                    time.mktime(
                        time.strptime(
                            response_json["createTime"].replace("Z", "+00:00"),
                            "%Y-%m-%dT%H:%M:%S.%f%z",
                        )
                    )
                ),
                filename=response_json.get("displayName", ""),
                object="file",
                purpose="user_data",
                status=status,
                status_details=(
                    str(response_json.get("error", ""))
                    if gemini_state == "FAILED"
                    else None
                ),
            )
        except Exception as e:
            verbose_logger.exception(f"Error parsing file retrieve response: {str(e)}")
            raise ValueError(f"Error parsing file retrieve response: {str(e)}")

    def transform_delete_file_request(
        self,
        file_id: str,
        optional_params: dict,
        litellm_params: dict,
    ) -> tuple[str, dict]:
        """
        Transform delete file request for Google AI Studio.

        Args:
            file_id: The file URI (e.g., "files/abc123" or full URI)
            optional_params: Optional parameters
            litellm_params: LiteLLM parameters containing api_key

        Returns:
            tuple[str, dict]: (url, params) for the DELETE request
        """
        api_base = self.get_api_base(litellm_params.get("api_base"))
        if not api_base:
            raise ValueError("api_base is required")

        # Get API key from multiple sources (same pattern as get_complete_url)
        api_key = litellm_params.get("api_key") or self.get_api_key()
        if not api_key:
            raise ValueError("api_key is required")

        # Normalize and encode the file name before interpolating it into the URL.
        file_name = self._normalize_gemini_file_id(file_id)

        # Construct the delete URL
        url = f"{api_base}/v1beta/{file_name}"

        # Add API key as header (Google AI Studio uses x-goog-api-key header)
        params: dict = {}

        return url, params

    def transform_delete_file_response(
        self,
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
    ) -> FileDeleted:
        """
        Transform Gemini's file delete response into OpenAI-style FileDeleted.

        Google AI Studio returns an empty JSON object {} on successful deletion.
        """
        try:
            # Google AI Studio returns {} on successful deletion
            if raw_response.status_code == 200:
                # Extract file ID from the request URL if possible
                file_id = "deleted"
                if hasattr(raw_response, "request") and raw_response.request:
                    url = str(raw_response.request.url)
                    if "/files/" in url:
                        file_id = url.split("/files/")[-1].split("?")[0]
                        # Add the files/ prefix if not present
                        if not file_id.startswith("files/"):
                            file_id = f"files/{file_id}"

                return FileDeleted(id=file_id, deleted=True, object="file")
            else:
                raise ValueError(f"Failed to delete file: {raw_response.text}")
        except Exception as e:
            verbose_logger.exception(f"Error parsing file delete response: {str(e)}")
            raise ValueError(f"Error parsing file delete response: {str(e)}")

    def transform_list_files_request(
        self,
        purpose: Optional[str],
        optional_params: dict,
        litellm_params: dict,
    ) -> tuple[str, dict]:
        raise NotImplementedError(
            "GoogleAIStudioFilesHandler does not support file listing"
        )

    def transform_list_files_response(
        self,
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
    ) -> List[OpenAIFileObject]:
        raise NotImplementedError(
            "GoogleAIStudioFilesHandler does not support file listing"
        )

    def transform_file_content_request(
        self,
        file_content_request,
        optional_params: dict,
        litellm_params: dict,
    ) -> tuple[str, dict]:
        raise NotImplementedError(
            "GoogleAIStudioFilesHandler does not support file content retrieval"
        )

    def transform_file_content_response(
        self,
        raw_response: httpx.Response,
        logging_obj: LiteLLMLoggingObj,
        litellm_params: dict,
    ) -> HttpxBinaryResponseContent:
        raise NotImplementedError(
            "GoogleAIStudioFilesHandler does not support file content retrieval"
        )
