# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import asyncio
import logging
import mimetypes
from collections.abc import Callable
from pathlib import Path
from typing import Any, Literal

import aiofiles
import httpx

from pyrit.common.deprecation import print_deprecation_message
from pyrit.models import (
    Message,
    MessagePiece,
    construct_response_from_request,
)
from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
from pyrit.prompt_target.common.target_configuration import TargetConfiguration
from pyrit.prompt_target.common.utils import limit_requests_per_minute
from pyrit.prompt_target.http_target.http_target import HTTPTarget

logger = logging.getLogger(__name__)


class HTTPXAPITarget(HTTPTarget):
    """
    A subclass of HTTPTarget that *only* does "API mode" (no raw HTTP request).
    This is a simpler approach for uploading files or sending JSON/form data.

    Additionally, if 'file_path' is not provided in the constructor,
    we attempt to pull it from the prompt's `converted_value`, assuming
    it's a local file path generated by a Converter (like PDFConverter).
    """

    _PATH_TYPES: frozenset[str] = frozenset({"image_path", "audio_path", "video_path", "binary_path"})
    _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration(
        capabilities=TargetCapabilities(
            supports_multi_turn=True,
            input_modalities=frozenset(
                {
                    frozenset(["text"]),
                    frozenset(["image_path"]),
                    frozenset(["audio_path"]),
                    frozenset(["video_path"]),
                    frozenset(["binary_path"]),
                }
            ),
        )
    )

    def __init__(
        self,
        *,
        http_url: str,
        method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"] = "POST",
        file_path: str | None = None,
        allowed_upload_directory: str | Path | None = None,
        json_data: dict[str, Any] | None = None,
        form_data: dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        http2: bool | None = None,
        follow_redirects: bool = True,
        callback_function: Callable[..., Any] | None = None,
        max_requests_per_minute: int | None = None,
        custom_configuration: TargetConfiguration | None = None,
        **httpx_client_kwargs: Any,
    ) -> None:
        """
        Force the parent 'HTTPTarget' to skip raw http_request logic by setting http_request=None.

        Args:
            http_url (str): The URL to send the HTTP request to.
            method (str): The HTTP method to use (GET, POST, PUT, DELETE, PATCH, HEAD, OPTIONS). Defaults to "POST".
            file_path (str, Optional): Path to a file to upload. If not provided, we attempt to pull it from the
            prompt's `converted_value`.
            allowed_upload_directory (str | Path | None): Directory containing files this target may upload.
                Required when uploading files. Defaults to None.
            json_data (dict, Optional): JSON data to send in the request body (for POST/PUT/PATCH).
            form_data (dict, Optional): Form data to send in the request body (for POST/PUT/PATCH).
            params (dict, Optional): Query parameters to include in the request URL (for GET/HEAD).
            headers (dict, Optional): Headers to include in the request.
            http2 (bool, Optional): Whether to use HTTP/2. If None, defaults to False.
            follow_redirects (bool): Whether to follow HTTP redirects. Defaults to True for backward compatibility;
                set to False when redirects are unnecessary or the destination must remain fixed.
            callback_function (Callable, Optional): Function to parse the HTTP response.
            max_requests_per_minute (int, Optional): Maximum number of requests per minute.
            custom_configuration (TargetConfiguration, Optional): Override the default configuration for this target
            instance.
            Defaults to None.
            **httpx_client_kwargs: Additional keyword arguments to pass to the httpx.AsyncClient constructor.

        Raises:
            ValueError: If the HTTP method is invalid.
            ValueError: If file uploads are attempted with an HTTP method that does not support them.
            ValueError: If the allowed upload directory does not exist or is not a directory.
        """
        super().__init__(
            http_request="",
            prompt_regex_string="",
            use_tls=True,
            follow_redirects=follow_redirects,
            callback_function=callback_function,
            max_requests_per_minute=max_requests_per_minute,
            custom_configuration=custom_configuration,
            **httpx_client_kwargs,
        )

        self.http_url = http_url
        self.method = method
        self.file_path = file_path
        self.allowed_upload_directory = self._resolve_allowed_upload_directory(allowed_upload_directory)
        self.json_data = json_data
        self.form_data = form_data
        self.params = params
        self.headers = headers or {}
        self.http2 = http2

        # Validate method
        if self.method not in {"GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"}:
            raise ValueError(f"Invalid HTTP method: {self.method}")

    @limit_requests_per_minute
    async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]:
        """
        Override the parent's method to skip raw http_request usage,
        and do a standard "API mode" approach.

        - If file_path is set or we can deduce it from the message piece, we upload a file.
        - Otherwise, we send normal requests with JSON or form_data (if provided).

        Returns:
            list[Message]: A list containing the response object with generated text pieces.

        Raises:
            ValueError: If no `http_url` is provided.
            httpx.TimeoutException: If the request times out.
            httpx.RequestError: If the request fails.
            FileNotFoundError: If the specified file to upload is not found.
        """
        message = normalized_conversation[-1]
        message_piece: MessagePiece = message.message_pieces[0]
        upload_path = await self._get_upload_path_async(message_piece=message_piece)
        if upload_path and self.method not in {"POST", "PUT"}:
            raise ValueError(f"File uploads are not allowed with HTTP method: {self.method}")

        if not self.http_url:
            raise ValueError("No `http_url` provided for HTTPXApiTarget.")

        http2_version = self.http2 if self.http2 is not None else False

        async with httpx.AsyncClient(http2=http2_version, **self.httpx_client_kwargs) as client:
            try:
                if upload_path:
                    # Handle file upload (only for POST & PUT)
                    filename = upload_path.name
                    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"

                    async with aiofiles.open(upload_path, "rb") as fp:
                        file_bytes = await fp.read()

                    files = {"file": (filename, file_bytes, mime_type)}
                    logger.info(f"HTTPXApiTarget: uploading file={filename} via {self.method} to {self.http_url}")

                    response = await client.request(
                        method=self.method,
                        url=self.http_url,
                        headers=self.headers,
                        params=self.params,
                        files=files,
                        follow_redirects=self.follow_redirects,
                    )
                else:
                    # No file upload, handle based on HTTP method
                    logger.info(f"HTTPXApiTarget: sending {self.method} to {self.http_url} with possible JSON/form.")
                    response = await client.request(
                        method=self.method,
                        url=self.http_url,
                        headers=self.headers,
                        params=self.params,
                        json=self.json_data if self.method in {"POST", "PUT", "PATCH"} else None,
                        data=self.form_data if self.method in {"POST", "PUT", "PATCH"} else None,
                        follow_redirects=self.follow_redirects,
                    )

            except httpx.TimeoutException:
                logger.error(f"Timeout error for URL: {self.http_url}")
                raise
            except httpx.RequestError as e:
                logger.error(f"Request failed: {e}")
                raise
            except FileNotFoundError as e:
                logger.error(f"File not found: {self.file_path}. Exception: {e}")
                raise

        response_content = response.content

        # If a callback function was set, let them parse the response
        if self.callback_function:
            response_content = self.callback_function(response=response)

        # Reuse parent's response object construction
        response_entry = construct_response_from_request(
            request=message_piece, response_text_pieces=[str(response_content)]
        )

        return [response_entry]

    async def _get_upload_path_async(self, *, message_piece: MessagePiece) -> Path | None:
        if self.file_path:
            candidate = Path(self.file_path)
        elif message_piece.converted_value_data_type in self._PATH_TYPES:
            candidate = Path(message_piece.converted_value)
        elif await asyncio.to_thread(Path(message_piece.converted_value).is_file):
            candidate = Path(message_piece.converted_value)
            print_deprecation_message(
                old_item="HTTPXAPITarget implicit text-path uploads",
                new_item="a path-typed MessagePiece or explicit file_path",
                removed_in="1.3.0",
            )
        else:
            return None

        if self.allowed_upload_directory is None:
            print_deprecation_message(
                old_item="HTTPXAPITarget file uploads without allowed_upload_directory",
                new_item="HTTPXAPITarget(..., allowed_upload_directory=...)",
                removed_in="1.3.0",
            )
            candidate = await asyncio.to_thread(candidate.resolve)
        else:
            candidate = await asyncio.to_thread(
                self._validate_upload_path,
                path=candidate,
                allowed_directory=self.allowed_upload_directory,
            )
        if not await asyncio.to_thread(candidate.is_file):
            raise FileNotFoundError(f"File not found: {candidate}")
        return candidate

    def _validate_upload_path(self, *, path: Path, allowed_directory: Path) -> Path:
        resolved_path = path.resolve()
        try:
            resolved_path.relative_to(allowed_directory)
        except ValueError as exc:
            raise ValueError(f"File upload path is outside the allowed upload directory: {resolved_path}") from exc
        return resolved_path

    @staticmethod
    def _resolve_allowed_upload_directory(allowed_directory: str | Path | None) -> Path | None:
        if allowed_directory is None:
            return None

        resolved_directory = Path(allowed_directory).resolve()
        if not resolved_directory.is_dir():
            raise ValueError(f"Allowed upload directory does not exist or is not a directory: {allowed_directory}")
        return resolved_directory