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

import logging
import mimetypes
import os
from typing import Any, Callable, Literal, Optional

import httpx

from pyrit.models import (
    PromptRequestPiece,
    PromptRequestResponse,
    construct_response_from_request,
)
from pyrit.prompt_target import HTTPTarget
from pyrit.prompt_target.common.utils import limit_requests_per_minute

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 PromptConverter (like PDFConverter).
    """

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

        self.http_url = http_url
        self.method = method
        self.file_path = file_path
        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}")

        # Validate file uploads (only `POST` and `PUT` allow file uploads)
        if self.file_path and self.method not in {"POST", "PUT"}:
            raise ValueError(f"File uploads are not allowed with HTTP method: {self.method}")

    @limit_requests_per_minute
    async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse:
        """
        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 prompt piece, we upload a file.
        - Otherwise, we send normal requests with JSON or form_data (if provided).
        """

        self._validate_request(prompt_request=prompt_request)
        request_piece: PromptRequestPiece = prompt_request.request_pieces[0]

        # If user didn't set file_path, see if the PDF path is in converted_value
        if not self.file_path:
            possible_path = request_piece.converted_value
            if isinstance(possible_path, str) and os.path.exists(possible_path):
                logger.info(f"HTTPXApiTarget: auto-using file_path from {possible_path}")
                self.file_path = possible_path

        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 self.file_path and os.path.exists(self.file_path):
                    # Handle file upload (only for POST & PUT)
                    filename = os.path.basename(self.file_path)
                    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"

                    with open(self.file_path, "rb") as fp:
                        file_bytes = 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,
                        files=files,
                        follow_redirects=True,
                    )
                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 if self.method in {"GET", "HEAD"} else None,
                        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=True,
                    )

            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=request_piece, response_text_pieces=[str(response_content)]
        )

        return response_entry
