import email.parser
import email.policy
import re
import threading
import shutil
import subprocess
import typing as t

import mercurial.hgweb.hgweb_mod_inner as web_inner
import mercurial.hgweb.request as hgreq
import mercurial.ui as hgui
from mercurial import extensions
from mercurial import wireprotoserver

type PermissionCheck = t.Callable[
    [web_inner.requestcontext, hgreq.parsedrequest, bytes],
    None,
]


CGI_VAR = re.compile(rb'[A-Z0-9_]+$')


def build_git_environ(
    req_ctx: web_inner.requestcontext,
    request: hgreq.parsedrequest,
) -> dict[bytes, bytes]:
    """Builds the environment to be sent to Git to serve HTTP."""
    fixed = {
        k: v for (k, v) in request.rawenv.items()
        if isinstance(v, bytes) and CGI_VAR.match(k)
    }
    fixed[b'GIT_HTTP_EXPORT_ALL'] = b'yes'
    fixed[b'GIT_PROJECT_ROOT'] = req_ctx.repo.path
    fixed[b'PATH_INFO'] = b'/git/' + request.dispatchpath
    return fixed


def parse_cgi_response(output: t.IO[bytes]) -> t.Tuple[bytes, dict[bytes, bytes], t.IO[bytes]]:
    parser = email.parser.BytesFeedParser(policy=email.policy.HTTP)
    while line := output.readline():
        if not line.rstrip(b'\r\n'):
            # We've reached the end of the headers.
            # Leave the rest in the output for later.
            break
        parser.feed(line)
    msg = parser.close()
    status = msg.get('Status', '200 OK I guess').encode('utf-8')
    del msg['Status']  # this won't raise an exception
    byte_headers = {k.encode('utf-8'): v.encode('utf-8') for (k, v) in msg.items()}
    return status, byte_headers, output


_feeds = 0
_feeder_lock = threading.Lock()


def feed_count() -> int:
    global _feeds
    with _feeder_lock:
        _feeds += 1
        return _feeds


def handle_git_protocol(
    original: t.Callable[..., bool],
    req_ctx: web_inner.requestcontext,
    request: hgreq.parsedrequest,
    response: hgreq.wsgiresponse,
    check_permission: PermissionCheck,
) -> bool:
    if request.headers.get(b'git-protocol'):
        # If a request is git, we assume we should be the one handling it.
        cgi_env = build_git_environ(req_ctx, request)
        http_backend = req_ctx.repo.ui.configlist(
            b'git-serve', b'http-backend', default=(b'git', b'http-backend'))
        is_post = request.method == b'POST'
        call = subprocess.Popen(
            http_backend,
            close_fds=True,
            stdin=subprocess.PIPE if is_post else None,
            stdout=subprocess.PIPE,
            env=cgi_env,
            text=False,
        )
        def feed():
            try:
                with call.stdin as stdin:
                   shutil.copyfileobj(request.bodyfh, stdin)
            except (OSError, BrokenPipeError):
                pass  # Expected; this just means it's closed.
        if is_post:
            threading.Thread(target=feed, name=f'git-feeder-{feed_count()}').start()
        status, headers, rest = parse_cgi_response(call.stdout)
        response.status = status
        for k, v in headers.items():
            response.headers[k] = v
        def write_the_rest():
            with call, rest:
                while more := rest.read(64 * 1024):
                    yield more
        response.setbodygen(write_the_rest())
        response.sendresponse()
        return True
    return original(req_ctx, request, response, check_permission)

def uisetup(_: hgui.ui) -> None:
    extensions.wrapfunction(wireprotoserver, 'handlewsgirequest', handle_git_protocol)