# --------------------------------------------------------------------------
#
# Copyright (c) Microsoft Corporation. All rights reserved.
#
# The MIT License (MIT)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the ""Software""), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# --------------------------------------------------------------------------

import json
import requests

try:
    from io import BytesIO
except ImportError:
    from cStringIO import StringIO as BytesIO
import xml.etree.ElementTree as ET
import sys

import pytest

from azure.core.configuration import Configuration
from azure.core.pipeline import Pipeline
from azure.core.rest import HttpRequest
from azure.core import PipelineClient
from azure.core.pipeline.policies import (
    SansIOHTTPPolicy,
    UserAgentPolicy,
    DistributedTracingPolicy,
    RedirectPolicy,
    RetryPolicy,
    HttpLoggingPolicy,
    HTTPPolicy,
    SansIOHTTPPolicy,
    SensitiveHeaderCleanupPolicy,
)
from azure.core.pipeline.transport._base import PipelineClientBase, _format_url_section
from azure.core.pipeline.transport import (
    HttpTransport,
    RequestsTransport,
)
from utils import HTTP_REQUESTS, is_rest

from azure.core.exceptions import AzureError
from azure.core.pipeline._base import cleanup_kwargs_for_transport


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_default_http_logging_policy(http_request):
    config = Configuration()
    pipeline_client = PipelineClient(base_url="test")
    pipeline = pipeline_client._build_pipeline(config)
    http_logging_policy = pipeline._impl_policies[-1]._policy
    assert http_logging_policy.allowed_header_names == HttpLoggingPolicy.DEFAULT_HEADERS_WHITELIST
    assert http_logging_policy.allowed_header_names == HttpLoggingPolicy.DEFAULT_HEADERS_ALLOWLIST
    assert "WWW-Authenticate" in http_logging_policy.allowed_header_names
    assert "x-vss-e2eid" in http_logging_policy.allowed_header_names
    assert "x-msedge-ref" in http_logging_policy.allowed_header_names
    # Testing I can replace the set entirely
    HttpLoggingPolicy.DEFAULT_HEADERS_ALLOWLIST = set(HttpLoggingPolicy.DEFAULT_HEADERS_ALLOWLIST)
    HttpLoggingPolicy.DEFAULT_HEADERS_WHITELIST = set(HttpLoggingPolicy.DEFAULT_HEADERS_ALLOWLIST)


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_pass_in_http_logging_policy(http_request):
    config = Configuration()
    http_logging_policy = HttpLoggingPolicy()
    http_logging_policy.allowed_header_names.update({"x-ms-added-header"})
    config.http_logging_policy = http_logging_policy

    pipeline_client = PipelineClient(base_url="test")
    pipeline = pipeline_client._build_pipeline(config)
    http_logging_policy = pipeline._impl_policies[-1]._policy
    assert http_logging_policy.allowed_header_names == HttpLoggingPolicy.DEFAULT_HEADERS_WHITELIST.union(
        {"x-ms-added-header"}
    )
    assert http_logging_policy.allowed_header_names == HttpLoggingPolicy.DEFAULT_HEADERS_ALLOWLIST.union(
        {"x-ms-added-header"}
    )


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_sans_io_exception(http_request):
    class BrokenSender(HttpTransport):
        def send(self, request, **config):
            raise ValueError("Broken")

        def open(self):
            self.session = requests.Session()

        def close(self):
            self.session.close()

        def __exit__(self, exc_type, exc_value, traceback):
            """Raise any exception triggered within the runtime context."""
            return self.close()

    pipeline = Pipeline(BrokenSender(), [SansIOHTTPPolicy()])

    req = http_request("GET", "/")
    with pytest.raises(ValueError):
        pipeline.run(req)

    class SwapExec(SansIOHTTPPolicy):
        def on_exception(self, requests, **kwargs):
            exc_type, exc_value, exc_traceback = sys.exc_info()
            raise NotImplementedError(exc_value)

    pipeline = Pipeline(BrokenSender(), [SwapExec()])
    with pytest.raises(NotImplementedError):
        pipeline.run(req)


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_requests_socket_timeout(http_request):
    conf = Configuration()
    request = http_request("GET", "https://bing.com")
    policies = [UserAgentPolicy("myusergant"), RedirectPolicy()]
    # Sometimes this will raise a read timeout, sometimes a socket timeout depending on timing.
    # Either way, the error should always be wrapped as an AzureError to ensure it's caught
    # by the retry policy.
    with pytest.raises(AzureError):
        with Pipeline(RequestsTransport(), policies=policies) as pipeline:
            response = pipeline.run(request, connection_timeout=0.000001, read_timeout=0.000001)


def test_format_url_basic():
    client = PipelineClientBase("https://bing.com")
    formatted = client.format_url("/{foo}", foo="bar")
    assert formatted == "https://bing.com/bar"


def test_format_url_with_query():
    client = PipelineClientBase("https://bing.com/path?query=testvalue&x=2ndvalue")
    formatted = client.format_url("/{foo}", foo="bar")
    assert formatted == "https://bing.com/path/bar?query=testvalue&x=2ndvalue"


def test_format_url_missing_param_values():
    client = PipelineClientBase("https://bing.com/path")
    formatted = client.format_url("/{foo}")
    assert formatted == "https://bing.com/path"


def test_format_url_missing_param_values_with_query():
    client = PipelineClientBase("https://bing.com/path?query=testvalue&x=2ndvalue")
    formatted = client.format_url("/{foo}")
    assert formatted == "https://bing.com/path?query=testvalue&x=2ndvalue"


def test_format_url_extra_path():
    client = PipelineClientBase("https://bing.com/path")
    formatted = client.format_url("/subpath/{foo}", foo="bar")
    assert formatted == "https://bing.com/path/subpath/bar"


def test_format_url_complex_params():
    client = PipelineClientBase("https://bing.com/path")
    formatted = client.format_url("/subpath/{a}/{b}/foo/{c}/bar", a="X", c="Y")
    assert formatted == "https://bing.com/path/subpath/X/foo/Y/bar"


def test_format_url_extra_path_missing_values():
    client = PipelineClientBase("https://bing.com/path")
    formatted = client.format_url("/subpath/{foo}")
    assert formatted == "https://bing.com/path/subpath"


def test_format_url_extra_path_missing_values_with_query():
    client = PipelineClientBase("https://bing.com/path?query=testvalue&x=2ndvalue")
    formatted = client.format_url("/subpath/{foo}")
    assert formatted == "https://bing.com/path/subpath?query=testvalue&x=2ndvalue"


def test_format_url_full_url():
    client = PipelineClientBase("https://bing.com/path")
    formatted = client.format_url("https://google.com/subpath/{foo}", foo="bar")
    assert formatted == "https://google.com/subpath/bar"


def test_format_url_no_base_url():
    client = PipelineClientBase(None)
    formatted = client.format_url("https://google.com/subpath/{foo}", foo="bar")
    assert formatted == "https://google.com/subpath/bar"


def test_format_url_double_query():
    client = PipelineClientBase("https://bing.com/path?query=testvalue&x=2ndvalue")
    formatted = client.format_url("/subpath?a=X&c=Y")
    assert formatted == "https://bing.com/path/subpath?query=testvalue&x=2ndvalue&a=X&c=Y"


def test_format_url_query_strings():
    client = PipelineClientBase("https://foo.core.windows.net")
    formatted = client.format_url("/")
    assert formatted == "https://foo.core.windows.net/"

    formatted = client.format_url("/?a=X&c=Y")
    assert formatted == "https://foo.core.windows.net/?a=X&c=Y"

    formatted = client.format_url("?a=X&c=Y")
    assert formatted == "https://foo.core.windows.net?a=X&c=Y"

    formatted = client.format_url("/Tables/?a=X&c=Y")
    assert formatted == "https://foo.core.windows.net/Tables/?a=X&c=Y"

    formatted = client.format_url("/Tables?a=X&c=Y")
    assert formatted == "https://foo.core.windows.net/Tables?a=X&c=Y"


def test_format_url_trailing_slash_preserved_with_query_only():
    # Test that trailing slash in base URL is preserved when url_template is query-string only
    # https://github.com/Azure/azure-sdk-for-python/issues/45365
    client = PipelineClientBase("{url}")
    formatted = client.format_url("?versionid=2026-02-25", url="https://storage.blob.core.windows.net/sample//a/a/")
    assert formatted == "https://storage.blob.core.windows.net/sample//a/a/?versionid=2026-02-25"


def test_format_url_from_http_request():
    client = PipelineClientBase("https://foo.core.windows.net")

    _url = "/"
    _params = {"foo": "bar"}
    request = HttpRequest("GET", _url, params=_params)
    formatted = client.format_url(request.url)
    assert formatted == "https://foo.core.windows.net/?foo=bar"

    _url = "?restype=service&comp=properties"
    request = HttpRequest("GET", _url)
    formatted = client.format_url(request.url)
    assert formatted == "https://foo.core.windows.net?restype=service&comp=properties"


def test_format_url_braces_with_dot():
    base_url = "https://bing.com/{aaa.bbb}"
    with pytest.raises(ValueError):
        url = _format_url_section(base_url)


def test_format_url_single_brace():
    base_url = "https://bing.com/{aaa.bbb"
    with pytest.raises(ValueError):
        url = _format_url_section(base_url)


def test_format_incorrect_endpoint():
    # https://github.com/Azure/azure-sdk-for-python/pull/12106
    client = PipelineClientBase("{Endpoint}/text/analytics/v3.0")
    with pytest.raises(ValueError) as exp:
        client.format_url("foo/bar")
    assert (
        str(exp.value) == "The value provided for the url part Endpoint was incorrect, and resulted in an invalid url"
    )


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_json(http_request):

    request = http_request("GET", "/")
    data = "Lots of dataaaa"
    request.set_json_body(data)

    assert request.data == json.dumps(data)
    assert request.headers.get("Content-Length") == "17"


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_data(http_request):

    request = http_request("GET", "/")
    data = "Lots of dataaaa"
    request.set_bytes_body(data)

    assert request.data == data
    assert request.headers.get("Content-Length") == "15"


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_stream(http_request):
    request = http_request("GET", "/")

    data = b"Lots of dataaaa"
    request.set_streamed_data_body(data)
    assert request.data == data

    def data_gen():
        for i in range(10):
            yield i

    data = data_gen()
    request.set_streamed_data_body(data)
    assert request.data == data

    data = BytesIO(b"Lots of dataaaa")
    request.set_streamed_data_body(data)
    assert request.data == data


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_xml(http_request):
    request = http_request("GET", "/")
    data = ET.Element("root")
    request.set_xml_body(data)

    assert request.data == b"<?xml version='1.0' encoding='utf-8'?>\n<root />"


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_url_with_params(http_request):

    request = http_request("GET", "/")
    request.url = "a/b/c?t=y"
    request.format_parameters({"g": "h"})

    assert request.url in ["a/b/c?g=h&t=y", "a/b/c?t=y&g=h"]


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_url_with_params_as_list(http_request):

    request = http_request("GET", "/")
    request.url = "a/b/c?t=y"
    request.format_parameters({"g": ["h", "i"]})

    assert request.url in ["a/b/c?g=h&g=i&t=y", "a/b/c?t=y&g=h&g=i"]


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_url_with_params_with_none_in_list(http_request):

    request = http_request("GET", "/")
    request.url = "a/b/c?t=y"
    with pytest.raises(ValueError):
        request.format_parameters({"g": ["h", None]})


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_url_with_params_with_none(http_request):

    request = http_request("GET", "/")
    request.url = "a/b/c?t=y"
    with pytest.raises(ValueError):
        request.format_parameters({"g": None})


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_repr(http_request):
    request = http_request("GET", "hello.com")
    assert repr(request) == "<HttpRequest [GET], url: 'hello.com'>"


def test_add_custom_policy():
    class BooPolicy(HTTPPolicy):
        def send(*args):
            raise AzureError("boo")

    class FooPolicy(HTTPPolicy):
        def send(*args):
            raise AzureError("boo")

    config = Configuration()
    retry_policy = RetryPolicy()
    config.retry_policy = retry_policy
    boo_policy = BooPolicy()
    foo_policy = FooPolicy()
    client = PipelineClient(base_url="test", config=config, per_call_policies=boo_policy)
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo < pos_retry

    client = PipelineClient(base_url="test", config=config, per_call_policies=[boo_policy])
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo < pos_retry

    client = PipelineClient(base_url="test", config=config, per_retry_policies=boo_policy)
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo > pos_retry

    client = PipelineClient(base_url="test", config=config, per_retry_policies=[boo_policy])
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo > pos_retry

    client = PipelineClient(base_url="test", config=config, per_call_policies=boo_policy, per_retry_policies=foo_policy)
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    assert foo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_foo = policies.index(foo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo < pos_retry
    assert pos_foo > pos_retry

    client = PipelineClient(
        base_url="test", config=config, per_call_policies=[boo_policy], per_retry_policies=[foo_policy]
    )
    policies = client._pipeline._impl_policies
    assert boo_policy in policies
    assert foo_policy in policies
    pos_boo = policies.index(boo_policy)
    pos_foo = policies.index(foo_policy)
    pos_retry = policies.index(retry_policy)
    assert pos_boo < pos_retry
    assert pos_foo > pos_retry

    policies = [UserAgentPolicy(), RetryPolicy(), DistributedTracingPolicy()]
    client = PipelineClient(base_url="test", policies=policies, per_call_policies=boo_policy)
    actual_policies = client._pipeline._impl_policies
    assert boo_policy == actual_policies[0]
    client = PipelineClient(base_url="test", policies=policies, per_call_policies=[boo_policy])
    actual_policies = client._pipeline._impl_policies
    assert boo_policy == actual_policies[0]

    client = PipelineClient(base_url="test", policies=policies, per_retry_policies=foo_policy)
    actual_policies = client._pipeline._impl_policies
    assert foo_policy == actual_policies[2]
    client = PipelineClient(base_url="test", policies=policies, per_retry_policies=[foo_policy])
    actual_policies = client._pipeline._impl_policies
    assert foo_policy == actual_policies[2]

    client = PipelineClient(
        base_url="test", policies=policies, per_call_policies=boo_policy, per_retry_policies=foo_policy
    )
    actual_policies = client._pipeline._impl_policies
    assert boo_policy == actual_policies[0]
    assert foo_policy == actual_policies[3]
    client = PipelineClient(
        base_url="test", policies=policies, per_call_policies=[boo_policy], per_retry_policies=[foo_policy]
    )
    actual_policies = client._pipeline._impl_policies
    assert boo_policy == actual_policies[0]
    assert foo_policy == actual_policies[3]

    policies = [UserAgentPolicy(), DistributedTracingPolicy()]
    with pytest.raises(ValueError):
        client = PipelineClient(base_url="test", policies=policies, per_retry_policies=foo_policy)
    with pytest.raises(ValueError):
        client = PipelineClient(base_url="test", policies=policies, per_retry_policies=[foo_policy])


def test_no_cleanup_policy_when_redirect_policy_is_empty():
    config = Configuration()
    client = PipelineClient(base_url="test", config=config)
    policies = client._pipeline._impl_policies
    for policy in policies:
        if isinstance(policy, SensitiveHeaderCleanupPolicy):
            assert False


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_basic_requests(port, http_request):

    conf = Configuration()
    request = http_request("GET", "http://localhost:{}/basic/string".format(port))
    policies = [UserAgentPolicy("myusergant"), RedirectPolicy()]
    with Pipeline(RequestsTransport(), policies=policies) as pipeline:
        response = pipeline.run(request)
        if is_rest(request):
            assert is_rest(response.http_response)

    assert pipeline._transport.session is None
    assert isinstance(response.http_response.status_code, int)


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_basic_options_requests(port, http_request):

    request = http_request("OPTIONS", "http://localhost:{}/basic/string".format(port))
    policies = [UserAgentPolicy("myusergant"), RedirectPolicy()]
    with Pipeline(RequestsTransport(), policies=policies) as pipeline:
        response = pipeline.run(request)
        if is_rest(request):
            assert is_rest(response.http_response)

    assert pipeline._transport.session is None
    assert isinstance(response.http_response.status_code, int)


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_basic_requests_separate_session(port, http_request):

    session = requests.Session()
    request = http_request("GET", "http://localhost:{}/basic/string".format(port))
    policies = [UserAgentPolicy("myusergant"), RedirectPolicy()]
    transport = RequestsTransport(session=session, session_owner=False)
    with Pipeline(transport, policies=policies) as pipeline:
        response = pipeline.run(request)
        if is_rest(request):
            assert is_rest(response.http_response)

    assert transport.session
    assert isinstance(response.http_response.status_code, int)
    transport.close()
    assert transport.session
    transport.session.close()


@pytest.mark.parametrize("http_request", HTTP_REQUESTS)
def test_request_text(port, http_request):
    client = PipelineClientBase("http://localhost:{}".format(port))
    if is_rest(http_request):
        request = http_request("GET", "/", json="foo")
    else:
        request = client.get("/", content="foo")

    # In absence of information, everything is JSON (double quote added)
    assert request.data == json.dumps("foo")

    if is_rest(http_request):
        request = http_request("POST", "/", headers={"content-type": "text/whatever"}, content="foo")
    else:
        request = client.post("/", headers={"content-type": "text/whatever"}, content="foo")

    # We want a direct string
    assert request.data == "foo"


def test_cleanup_kwargs():
    kwargs = {"insecure_domain_change": True, "enable_cae": True}
    cleanup_kwargs_for_transport(kwargs)
    assert "insecure_domain_change" not in kwargs
    assert "enable_cae" not in kwargs
