import subprocess

from httpx import Request, Response, HTTPTransport, AsyncHTTPTransport

def resolve_host(host, dns_server="8.8.8.8") -> list[str]:
    try:
        result = subprocess.run(
            ["dig", "+short", host, f"@{dns_server}"],
            capture_output=True,
            text=True,
            check=True
        )
        ips = result.stdout.strip().split("\n")
        return [ip for ip in ips if ip]
    except subprocess.CalledProcessError:
        return []

class NameSolver:
    def get(self, name: str) -> str:
        if name == 'clients.plex.tv':
            hosts = resolve_host(name)
            if hosts:
                return hosts[0]

        return ''

    def resolve(self, request: Request) -> Request:
        host = request.url.host
        ip = self.get(host)

        if ip:
            request.extensions["sni_hostname"] = host
            request.url = request.url.copy_with(host=ip)

        return request


class CustomHost(HTTPTransport):
    def __init__(self, solver: NameSolver, *args, **kwargs) -> None:
        self.solver = solver
        super().__init__(*args, **kwargs)

    def handle_request(self, request: Request) -> Response:
        request = self.solver.resolve(request)
        return super().handle_request(request)


class AsyncCustomHost(AsyncHTTPTransport):
    def __init__(self, solver: NameSolver, *args, **kwargs) -> None:
        self.solver = solver
        super().__init__(*args, **kwargs)

    async def handle_async_request(self, request: Request) -> Response:
        request = self.solver.resolve(request)
        return await super().handle_async_request(request)