# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from enum import Enum
import logging
import os
from typing import TYPE_CHECKING, Optional
import urllib.parse as url_parse

from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
from azure.core.pipeline.policies import ContentDecodePolicy

# the functions we patch
try:
    from azure.core.pipeline.transport import RequestsTransport
except:
    pass

try:
    import httpx

    HTTPXTransport = httpx.HTTPTransport
    AsyncHTTPXTransport = httpx.AsyncHTTPTransport
except ImportError:
    httpx = None
    HTTPXTransport = None
    AsyncHTTPXTransport = None

from .config import PROXY_URL
from .helpers import (
    get_http_client,
    get_test_id,
    is_live,
    is_live_and_not_recording,
    set_recording_id,
    trim_kwargs_from_test_function,  # to clean up incoming arguments to the test function we are wrapping
)
from .proxy_startup import discovered_roots
import json

if TYPE_CHECKING:
    from typing import Callable, Dict, Tuple
    from azure.core.pipeline.transport import HttpRequest


# defaults
RECORDING_START_URL = "{}/record/start".format(PROXY_URL)
RECORDING_STOP_URL = "{}/record/stop".format(PROXY_URL)
PLAYBACK_START_URL = "{}/playback/start".format(PROXY_URL)
PLAYBACK_STOP_URL = "{}/playback/stop".format(PROXY_URL)


class RecordedTransport(str, Enum):
    """Enum for specifying which transports to record in the test proxy."""

    AZURE_CORE = "azure_core"
    HTTPX = "httpx"


def get_recording_assets(test_id: str) -> Optional[str]:
    """Used to retrieve the assets.json given a PYTEST_CURRENT_TEST test id."""
    for root in discovered_roots:
        current_dir = os.path.dirname(test_id)
        while current_dir is not None and not (os.path.dirname(current_dir) == current_dir):
            possible_assets = os.path.join(current_dir, "assets.json")
            possible_root = os.path.join(current_dir, ".git")

            # we need to check for assets.json first!
            if os.path.exists(os.path.join(root, possible_assets)):
                complete_path = os.path.abspath(os.path.join(root, possible_assets))
                return os.path.relpath(complete_path, root).replace("\\", "/")
            # we need the git check to prevent ascending out of the repo
            elif os.path.exists(os.path.join(root, possible_root)):
                return None
            else:
                current_dir = os.path.dirname(current_dir)

    return None


def start_record_or_playback(test_id: str) -> "Tuple[str, Dict[str, str]]":
    """Sends a request to begin recording or playing back the provided test.

    This returns a tuple, (a, b), where a is the recording ID of the test and b is the `variables` dictionary that maps
    test variables to values. If no variable dictionary was stored when the test was recorded, b is an empty dictionary.
    """
    variables = {}  # this stores a dictionary of test variable values that could have been stored with a recording

    json_payload = {"x-recording-file": test_id}
    assets_json = get_recording_assets(test_id)
    if assets_json:
        json_payload["x-recording-assets-file"] = assets_json

    encoded_payload = json.dumps(json_payload).encode("utf-8")
    http_client = get_http_client()

    if is_live():
        result = http_client.request(
            method="POST",
            url=RECORDING_START_URL,
            body=encoded_payload,
        )
        if result.status != 200:
            raise HttpResponseError(message=result.data)
        recording_id = result.headers["x-recording-id"]

    else:
        result = http_client.request(
            method="POST",
            url=PLAYBACK_START_URL,
            body=encoded_payload,
        )
        if result.status != 200:
            raise HttpResponseError(message=result.data)

        try:
            recording_id = result.headers["x-recording-id"]
        except KeyError as ex:
            raise ValueError(f"No recording file found for {test_id}") from ex
        if result.data:
            try:
                variables = json.loads(result.data.decode("utf-8"))
            except ValueError as ex:  # would be a JSONDecodeError, which subclasses ValueError
                raise ValueError("The response body returned from starting playback did not contain valid JSON") from ex

    # set recording ID in a module-level variable so that sanitizers can access it
    set_recording_id(test_id, recording_id)
    return (recording_id, variables)


def stop_record_or_playback(test_id: str, recording_id: str, test_variables: "Dict[str, str]") -> None:
    http_client = get_http_client()
    if is_live():
        response = http_client.request(
            method="POST",
            url=RECORDING_STOP_URL,
            headers={
                "x-recording-file": test_id,
                "x-recording-id": recording_id,
                "x-recording-save": "true",
                "Content-Type": "application/json",
            },
            # tests don't record successfully unless test_variables is a dictionary
            body=json.dumps(test_variables or {}).encode("utf-8"),
        )
    else:
        response = http_client.request(
            method="POST",
            url=PLAYBACK_STOP_URL,
            headers={"x-recording-id": recording_id},
        )
    if response.status >= 400:
        raise HttpResponseError(
            "The test proxy ran into an error while ending the session. Make sure any test variables you record have "
            f"string values. Full error details: {response.data}"
        )


def get_proxy_netloc() -> "Dict[str, str]":
    parsed_result = url_parse.urlparse(PROXY_URL)
    return {"scheme": parsed_result.scheme, "netloc": parsed_result.netloc}


def transform_request(request: "HttpRequest", recording_id: str) -> None:
    """Redirect the request to the test proxy, and store the original request URI in a header"""
    headers = request.headers

    parsed_result = url_parse.urlparse(request.url)
    updated_target = parsed_result._replace(**get_proxy_netloc()).geturl()
    if headers.get("x-recording-upstream-base-uri", None) is None:
        headers["x-recording-upstream-base-uri"] = "{}://{}".format(parsed_result.scheme, parsed_result.netloc)
    headers["x-recording-id"] = recording_id
    headers["x-recording-mode"] = "record" if is_live() else "playback"
    request.url = updated_target


def transform_httpx_request(request, recording_id: str) -> None:
    """Transform an httpx.Request to route through the test proxy."""
    parsed_result = url_parse.urlparse(str(request.url))

    # Store original upstream URI
    if "x-recording-upstream-base-uri" not in request.headers:
        request.headers["x-recording-upstream-base-uri"] = f"{parsed_result.scheme}://{parsed_result.netloc}"

    # Set recording headers
    request.headers["x-recording-id"] = recording_id
    request.headers["x-recording-mode"] = "record" if is_live() else "playback"

    # Rewrite URL to proxy
    updated_target = parsed_result._replace(**get_proxy_netloc()).geturl()
    request.url = httpx.URL(updated_target)


def restore_httpx_response_url(response) -> None:
    """Restore the response's request URL to the original upstream target."""
    try:
        parsed_resp = url_parse.urlparse(str(response.request.url))
        upstream_uri_str = response.request.headers.get("x-recording-upstream-base-uri", "")
        if upstream_uri_str:
            upstream_uri = url_parse.urlparse(upstream_uri_str)
            original_target = parsed_resp._replace(
                scheme=upstream_uri.scheme or parsed_resp.scheme, netloc=upstream_uri.netloc
            ).geturl()
            response.request.url = httpx.URL(original_target)
    except Exception:
        # Best-effort restore; don't fail the call if something goes wrong
        pass


def _transform_args(recording_id: str, *call_args, **call_kwargs):
    """Transform azure.core transport call arguments to route through the test proxy.

    Used by both sync and async decorators.
    """
    copied_positional_args = list(call_args)
    request = copied_positional_args[1]
    transform_request(request, recording_id)
    return tuple(copied_positional_args), call_kwargs


def _transform_httpx_args(recording_id: str, *call_args, **call_kwargs):
    """Transform httpx transport call arguments to route through the test proxy.

    Used by both sync and async decorators.
    """
    copied_positional_args = list(call_args)
    request = copied_positional_args[1]
    transform_httpx_request(request, recording_id)
    return tuple(copied_positional_args), call_kwargs


def recorded_by_proxy(*transports):
    """
    Decorator for recording and playing back test proxy sessions.

    Args:
        *transports: Which transport(s) to record. Pass one or more comma separated RecordedTransport enum values.
            - No args (default): Record RequestsTransport.send calls (azure.core).
            - RecordedTransport.AZURE_CORE: Record RequestsTransport.send calls. Same as the default above.
            - RecordedTransport.HTTPX: Record HTTPXTransport.handle_request calls.
            - RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX: Record both transports.

    Usages:
      from devtools_testutils import recorded_by_proxy, RecordedTransport

      # If your test uses azure.core only network calls (default)
      @recorded_by_proxy
      def test(...): ...

      # Explicitly enable azure.core recordings only (equivalent to the above)
      @recorded_by_proxy(RecordedTransport.AZURE_CORE)
      def test(...): ...

      # If your test uses httpx only for network calls
      @recorded_by_proxy(RecordedTransport.HTTPX)
      def test(...): ...

      # If your test uses both azure.core and httpx for network calls
      @recorded_by_proxy(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
      def test(...): ...
    """

    # Bare decorator usage: @recorded_by_proxy
    if len(transports) == 1 and callable(transports[0]):
        test_func = transports[0]
        transport_list = [(RequestsTransport, "send")]
        return _make_proxy_decorator(transport_list)(test_func)

    # Parameterized decorator usage: @recorded_by_proxy(...)
    # Determine which transports to use
    transport_list = []

    # If no transports specified, default to azure.core
    transport_set = set(transports) if transports else {RecordedTransport.AZURE_CORE}

    # Add transports based on what's in the set
    for transport in transport_set:
        if transport == RecordedTransport.AZURE_CORE or (
            isinstance(transport, str) and transport == RecordedTransport.AZURE_CORE.value
        ):
            transport_list.append((RequestsTransport, "send"))
        elif transport == RecordedTransport.HTTPX or (
            isinstance(transport, str) and transport == RecordedTransport.HTTPX.value
        ):
            if HTTPXTransport is not None:
                transport_list.append((HTTPXTransport, "handle_request"))

    # If still no transports, fall back to azure.core
    if not transport_list:
        transport_list = [(RequestsTransport, "send")]

    # Return a decorator function that will be applied to the test function
    return lambda test_func: _make_proxy_decorator(transport_list)(test_func)


def _make_proxy_decorator(transports):
    def _decorator(test_func: "Callable"):
        def record_wrap(*args, **kwargs):
            # ---- your existing trimming/early-exit logic ----
            trimmed_kwargs = {k: v for k, v in kwargs.items()}
            trim_kwargs_from_test_function(test_func, trimmed_kwargs)

            if is_live_and_not_recording():
                return test_func(*args, **trimmed_kwargs)

            test_id = get_test_id()
            recording_id, variables = start_record_or_playback(test_id)

            # Build a wrapper factory so each patched method closes over its own original
            def make_combined_call(original_transport_func, is_httpx=False):
                def combined_call(*call_args, **call_kwargs):
                    if is_httpx:
                        adjusted_args, adjusted_kwargs = _transform_httpx_args(recording_id, *call_args, **call_kwargs)
                        result = original_transport_func(*adjusted_args, **adjusted_kwargs)
                        restore_httpx_response_url(result)
                    else:
                        adjusted_args, adjusted_kwargs = _transform_args(recording_id, *call_args, **call_kwargs)
                        result = original_transport_func(*adjusted_args, **adjusted_kwargs)
                        # rewrite request.url to the original upstream for LROs, etc.
                        parsed_result = url_parse.urlparse(result.request.url)
                        upstream_uri = url_parse.urlparse(result.request.headers["x-recording-upstream-base-uri"])
                        upstream_uri_dict = {"scheme": upstream_uri.scheme, "netloc": upstream_uri.netloc}
                        original_target = parsed_result._replace(**upstream_uri_dict).geturl()
                        result.request.url = original_target
                    return result

                return combined_call

            # Patch multiple transports and ensure restoration
            test_variables = None
            test_run = False
            originals = []
            # monkeypatch all requested transports
            for owner, name in transports:
                original = getattr(owner, name)
                # Check if this is an httpx transport by comparing with httpx transport classes
                is_httpx_transport = (
                    (HTTPXTransport is not None and owner is HTTPXTransport)
                    or (AsyncHTTPXTransport is not None and owner is AsyncHTTPXTransport)
                    or (httpx is not None and owner.__module__.startswith("httpx"))
                )
                setattr(owner, name, make_combined_call(original, is_httpx=is_httpx_transport))
                originals.append((owner, name, original))

            try:
                try:
                    test_variables = test_func(*args, variables=variables, **trimmed_kwargs)
                    test_run = True
                except TypeError as error:
                    if "unexpected keyword argument" in str(error) and "variables" in str(error):
                        logger = logging.getLogger()
                        logger.info(
                            "This test can't accept variables as input. "
                            "Accept `**kwargs` and/or a `variables` parameter to use recorded variables."
                        )
                    else:
                        raise

                if not test_run:
                    test_variables = test_func(*args, **trimmed_kwargs)

            except ResourceNotFoundError as error:
                error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
                troubleshoot = (
                    "Playback failure -- for help resolving, see https://aka.ms/azsdk/python/test-proxy/troubleshoot."
                )
                message = error_body.get("message") or error_body.get("Message")
                error_with_message = ResourceNotFoundError(
                    message=f"{troubleshoot} Error details:\n{message}",
                    response=error.response,
                )
                raise error_with_message from error

            finally:
                # restore in reverse order
                for owner, name, original in reversed(originals):
                    setattr(owner, name, original)
                stop_record_or_playback(test_id, recording_id, test_variables)

            return test_variables

        return record_wrap

    return _decorator