# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import hashlib
from datetime import timedelta
from http.server import BaseHTTPRequestHandler, HTTPServer
from threading import Thread
from urllib.parse import parse_qs, urlparse

import pytest
import requests

from pants.engine.download_file import URLDownloadHandler, download_file
from pants.engine.fs import (
    Digest,
    DigestEntries,
    DownloadFile,
    FileDigest,
    FileEntry,
    NativeDownloadFile,
)
from pants.engine.internals.scheduler import ExecutionError
from pants.engine.rules import QueryRule
from pants.engine.unions import UnionMembership, UnionRule
from pants.option.global_options import GlobalOptions
from pants.testutil.option_util import create_subsystem
from pants.testutil.rule_runner import RuleRunner, run_rule_with_mocks

DOWNLOADS_FILE_DIGEST = FileDigest(
    "8fcbc50cda241aee7238c71e87c27804e7abc60675974eaf6567aa16366bc105", 14
)
DOWNLOADS_EXPECTED_DIRECTORY_DIGEST = Digest(
    "4c9cf91fcd7ba1abbf7f9a0a1c8175556a82bee6a398e34db3284525ac24a3ad", 84
)


@pytest.fixture
def global_options() -> GlobalOptions:
    return create_subsystem(
        GlobalOptions, file_downloads_retry_delay=0.01, file_downloads_max_attempts=2
    )


def test_no_union_members(global_options: GlobalOptions) -> None:
    union_membership = UnionMembership.empty()
    digest = run_rule_with_mocks(
        download_file,
        rule_args=[
            DownloadFile("http://pantsbuild.com/file.txt", DOWNLOADS_FILE_DIGEST),
            union_membership,
            global_options,
        ],
        mock_calls={
            "pants.engine.intrinsics.download_file": lambda _: DOWNLOADS_EXPECTED_DIRECTORY_DIGEST,
        },
        union_membership=union_membership,
    )
    assert digest == DOWNLOADS_EXPECTED_DIRECTORY_DIGEST


@pytest.mark.parametrize(
    "scheme, authority, url",
    [
        # Anything (every URL matches)
        (None, None, "s3://pantsbuild.com/file.txt"),
        (None, None, "http://pantsbuild.com/file.txt"),
        (None, None, "http://awesome.pantsbuild.com/file.txt"),
        # Scheme
        ("s3", None, "s3://pantsbuild.com/file.txt"),
        ("http*", None, "http://pantsbuild.com/file.txt"),
        ("http*", None, "https://pantsbuild.com/file.txt"),
        # Authority
        (None, "pantsbuild.com", "s3://pantsbuild.com/file.txt"),
        (None, "pantsbuild.com", "http://pantsbuild.com/file.txt"),
        (None, "pantsbuild.com", "https://pantsbuild.com/file.txt"),
        (None, "*.pantsbuild.com", "https://awesome.pantsbuild.com/file.txt"),
        (None, "*.pantsbuild.com*", "https://awesome.pantsbuild.com/file.txt"),
        (None, "*.pantsbuild.com*", "https://awesome.pantsbuild.com:80/file.txt"),
        # Both
        ("http*", "*.pantsbuild.com", "http://awesome.pantsbuild.com/file.txt"),
    ],
)
def test_matches(scheme, authority, url, global_options: GlobalOptions) -> None:
    class UnionMember(URLDownloadHandler):
        match_scheme = scheme
        match_authority = authority

        def mock_rule(self) -> Digest:
            assert isinstance(self, UnionMember)
            return DOWNLOADS_EXPECTED_DIRECTORY_DIGEST

    union_membership = UnionMembership.from_rules({UnionRule(URLDownloadHandler, UnionMember)})

    digest = run_rule_with_mocks(
        download_file,
        rule_args=[
            DownloadFile(url, DOWNLOADS_FILE_DIGEST),
            union_membership,
            global_options,
        ],
        mock_calls={
            "pants.engine.download_file.download_file_using_handler": UnionMember.mock_rule,
        },
        union_membership=union_membership,
    )
    assert digest == DOWNLOADS_EXPECTED_DIRECTORY_DIGEST


@pytest.mark.parametrize(
    "scheme, authority, url",
    [
        # Scheme
        ("s3", None, "http://pantsbuild.com/file.txt"),
        ("s3", None, "as3://pantsbuild.com/file.txt"),
        ("http", None, "https://pantsbuild.com/file.txt"),
        # Authority
        (None, "pantsbuild.com", "http://pantsbuild.com:80/file.txt"),
        (None, "*.pantsbuild.com", "https://pantsbuild.com/file.txt"),
        # Both
        ("http", "*.pantsbuild.com", "https://awesome.pantsbuild.com/file.txt"),
        ("https", "*.pantsbuild.com", "https://pantsbuild.com/file.txt"),
    ],
)
def test_doesnt_match(scheme, authority, url, global_options: GlobalOptions) -> None:
    class UnionMember(URLDownloadHandler):
        match_scheme = scheme
        match_authority = authority

    union_membership = UnionMembership.from_rules({UnionRule(URLDownloadHandler, UnionMember)})

    digest = run_rule_with_mocks(
        download_file,
        rule_args=[
            DownloadFile(url, DOWNLOADS_FILE_DIGEST),
            union_membership,
            global_options,
        ],
        mock_calls={
            "pants.engine.intrinsics.download_file": lambda _: DOWNLOADS_EXPECTED_DIRECTORY_DIGEST,
        },
        union_membership=union_membership,
    )
    assert digest == DOWNLOADS_EXPECTED_DIRECTORY_DIGEST


def test_too_many_matches(global_options: GlobalOptions) -> None:
    class AuthorityMatcher(URLDownloadHandler):
        match_authority = "pantsbuild.com"

    class SchemeMatcher(URLDownloadHandler):
        match_scheme = "http"

    union_membership = UnionMembership.from_rules(
        {
            UnionRule(URLDownloadHandler, AuthorityMatcher),
            UnionRule(URLDownloadHandler, SchemeMatcher),
        }
    )

    with pytest.raises(Exception, match=r"Too many registered URL handlers"):
        run_rule_with_mocks(
            download_file,
            rule_args=[
                DownloadFile("http://pantsbuild.com/file.txt", DOWNLOADS_FILE_DIGEST),
                union_membership,
                global_options,
            ],
            mock_calls={
                "pants.engine.intrinsics.download_file": lambda _: DOWNLOADS_EXPECTED_DIRECTORY_DIGEST,
            },
            union_membership=union_membership,
        )


def test_query_string_included_in_cache_key() -> None:
    class RequestHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            parsed_url = urlparse(self.path)
            query_params = parse_qs(parsed_url.query)
            body = query_params["val"][0].encode()
            self.send_response(200)
            self.send_header("Content-length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

    server = HTTPServer(("", 0), RequestHandler)
    port = server.server_port

    def _http_server_thread() -> None:
        server.serve_forever()

    t = Thread(target=_http_server_thread)
    t.daemon = True
    t.start()

    # Wait until the http server is operational.
    wait_attempts_remaining = 4
    while wait_attempts_remaining > 0:
        r = requests.get(f"http://127.0.0.1:{port}/?val=test", timeout=0.1)
        if r.status_code == 200:
            break
        wait_attempts_remaining -= 1

    if wait_attempts_remaining == 0:
        raise Exception("HTTP server thread did not startup.")

    rule_runner = RuleRunner(
        rules=[QueryRule(DigestEntries, (Digest,))],
        isolated_local_store=True,
    )

    response = "world"
    expected_digest = FileDigest(hashlib.sha256(response.encode()).hexdigest(), len(response))

    digest = rule_runner.request(
        Digest,
        [
            NativeDownloadFile(
                f"http://127.0.0.1:{port}/hello?val={response}",
                expected_digest=expected_digest,
                retry_delay_duration=timedelta(milliseconds=1),
                max_attempts=2,
            )
        ],
    )

    entries = rule_runner.request(DigestEntries, [digest])
    assert len(entries) == 1
    entry = entries[0]
    assert isinstance(entry, FileEntry)
    assert entry.file_digest == expected_digest

    with pytest.raises(ExecutionError, match=r"Downloaded file was larger than expected digest"):
        # Note: In manual testing before the bug fix in the Rust intrinsic, this call would succeed but would
        # return the cached content from the prior call to the intrinsic, which is not what should
        # have occurred. Instead, the expected error should have occurred.
        _ = rule_runner.request(
            Digest,
            [
                NativeDownloadFile(
                    f"http://127.0.0.1:{port}/hello?val=galaxy",
                    expected_digest=expected_digest,
                    retry_delay_duration=timedelta(milliseconds=1),
                    max_attempts=2,
                )
            ],
        )
