import socket
import socketserver
import select
import urllib.parse

BUFFER_SIZE = 65536
TIMEOUT = 30


def connect_upstream(host, port):
    return socket.create_connection((host, port), timeout=TIMEOUT)


class ProxyHandler(socketserver.BaseRequestHandler):
    def handle(self):
        client = self.request
        client.settimeout(TIMEOUT)
        data = b''
        while b'\r\n\r\n' not in data and len(data) < 65536:
            chunk = client.recv(BUFFER_SIZE)
            if not chunk:
                return
            data += chunk
        header_blob, _, extra = data.partition(b'\r\n\r\n')
        lines = header_blob.decode('iso-8859-1', 'replace').split('\r\n')
        if not lines:
            return
        parts = lines[0].split()
        if len(parts) < 3:
            return
        method, target, version = parts[0].upper(), parts[1], parts[2]
        header_lines = lines[1:]

        if method == 'CONNECT':
            host, _, port_s = target.partition(':')
            port = int(port_s or 443)
            remote = None
            try:
                remote = connect_upstream(host, port)
                remote.sendall(extra)
                client.sendall(b'HTTP/1.1 200 Connection Established\r\n\r\n')
                self.relay(client, remote)
            except Exception as e:
                try:
                    msg = ('HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/plain\r\nConnection: close\r\n\r\n' + str(e)).encode()
                    client.sendall(msg)
                except Exception:
                    pass
                if remote:
                    remote.close()
            return

        url = urllib.parse.urlsplit(target)
        host = url.hostname
        port = url.port or (443 if url.scheme == 'https' else 80)
        if not host:
            return
        path = urllib.parse.urlunsplit(('', '', url.path or '/', url.query, ''))
        remote = None
        try:
            remote = connect_upstream(host, port)
            remote.sendall(f'{method} {path} {version}\r\n'.encode('iso-8859-1'))
            for h in header_lines:
                if not h.lower().startswith('proxy-connection:'):
                    remote.sendall((h + '\r\n').encode('iso-8859-1'))
            remote.sendall(b'\r\n')
            if extra:
                remote.sendall(extra)
            self.relay(client, remote)
        except Exception:
            if remote:
                remote.close()

    def relay(self, client, remote):
        sockets = [client, remote]
        try:
            while True:
                readable, _, _ = select.select(sockets, [], [], TIMEOUT)
                if not readable:
                    break
                for s in readable:
                    data = s.recv(BUFFER_SIZE)
                    if not data:
                        return
                    (remote if s is client else client).sendall(data)
        finally:
            remote.close()


class ThreadingTCPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
    allow_reuse_address = True
    daemon_threads = True


if __name__ == '__main__':
    with ThreadingTCPServer(('127.0.0.1', 18888), ProxyHandler) as server:
        print('Local proxy listening on 127.0.0.1:18888', flush=True)
        server.serve_forever()