"""
Local HTTP CONNECT proxy for Playwright — handles proxy auth that Chromium can't do natively.

Architecture: Playwright → localhost:PORT (no auth) → proxy.iploop.io:8880 (with auth)

Key fixes over v1:
- Fresh upstream connection per request (no stale connections)
- Proper bidirectional tunnel (no premature close)
- Graceful shutdown on both directions
"""

import asyncio
import socket
import base64
import logging
from typing import Optional

logger = logging.getLogger("iploop.render_proxy")


class AuthProxy:
    """Local proxy that injects Proxy-Authorization into CONNECT requests."""
    
    def __init__(self, api_key: str, upstream_host: str = "proxy.iploop.io", upstream_port: int = 8880):
        self.api_key = api_key
        self.upstream_host = upstream_host
        self.upstream_port = upstream_port
        self.server = None
        self.port = None
        self._active = 0
    
    async def start(self) -> int:
        """Start on a random port. Returns port number."""
        sock = socket.socket()
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('127.0.0.1', 0))
        self.port = sock.getsockname()[1]
        sock.close()
        
        self.server = await asyncio.start_server(
            self._handle_client, '127.0.0.1', self.port
        )
        logger.debug(f"Auth proxy started on :{self.port}")
        return self.port
    
    async def stop(self):
        if self.server:
            self.server.close()
            await self.server.wait_closed()
            logger.debug(f"Auth proxy stopped (:{self.port})")
    
    async def _handle_client(self, cr: asyncio.StreamReader, cw: asyncio.StreamWriter):
        """Handle one CONNECT request from Playwright."""
        self._active += 1
        req_id = self._active
        ur = uw = None
        
        try:
            # 1. Read CONNECT request line
            line = await asyncio.wait_for(cr.readline(), timeout=10)
            if not line:
                return
            request = line.decode('utf-8', errors='ignore').strip()
            
            if not request.startswith('CONNECT '):
                cw.write(b'HTTP/1.1 400 Bad Request\r\n\r\n')
                await cw.drain()
                return
            
            target = request.split(' ')[1]
            logger.debug(f"[{req_id}] CONNECT {target}")
            
            # 2. Read remaining client headers (discard)
            while True:
                h = await asyncio.wait_for(cr.readline(), timeout=5)
                if h == b'\r\n' or not h:
                    break
            
            # 3. Connect to upstream proxy
            try:
                ur, uw = await asyncio.wait_for(
                    asyncio.open_connection(self.upstream_host, self.upstream_port),
                    timeout=10
                )
            except Exception as e:
                logger.debug(f"[{req_id}] Upstream connect failed: {e}")
                cw.write(b'HTTP/1.1 502 Bad Gateway\r\n\r\n')
                await cw.drain()
                return
            
            # 4. Send CONNECT with auth to upstream
            auth = base64.b64encode(f"user:{self.api_key}".encode()).decode()
            uw.write(
                f"CONNECT {target} HTTP/1.1\r\n"
                f"Host: {target}\r\n"
                f"Proxy-Authorization: Basic {auth}\r\n"
                f"\r\n".encode()
            )
            await uw.drain()
            
            # 5. Read upstream response
            resp_buf = b''
            while b'\r\n\r\n' not in resp_buf:
                chunk = await asyncio.wait_for(ur.read(4096), timeout=15)
                if not chunk:
                    break
                resp_buf += chunk
            
            status_line = resp_buf.split(b'\r\n')[0].decode('utf-8', errors='ignore')
            
            if '200' not in status_line:
                logger.debug(f"[{req_id}] Upstream rejected: {status_line}")
                cw.write(b'HTTP/1.1 502 Bad Gateway\r\n\r\n')
                await cw.drain()
                return
            
            # 6. Tell client tunnel is established
            cw.write(b'HTTP/1.1 200 Connection Established\r\n\r\n')
            await cw.drain()
            
            # 7. Forward any extra data from upstream response
            header_end = resp_buf.find(b'\r\n\r\n') + 4
            if len(resp_buf) > header_end:
                cw.write(resp_buf[header_end:])
                await cw.drain()
            
            # 8. Bidirectional tunnel — run until either side closes
            async def pipe(reader, writer, label):
                try:
                    while True:
                        data = await reader.read(16384)
                        if not data:
                            break
                        writer.write(data)
                        await writer.drain()
                except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError):
                    pass
                except Exception as e:
                    logger.debug(f"[{req_id}] {label}: {e}")
            
            # Run both directions, cancel the other when one finishes
            t1 = asyncio.create_task(pipe(cr, uw, "c→u"))
            t2 = asyncio.create_task(pipe(ur, cw, "u→c"))
            
            done, pending = await asyncio.wait(
                {t1, t2}, return_when=asyncio.FIRST_COMPLETED
            )
            for t in pending:
                t.cancel()
                try:
                    await t
                except asyncio.CancelledError:
                    pass
            
            logger.debug(f"[{req_id}] Tunnel closed for {target}")
            
        except asyncio.TimeoutError:
            logger.debug(f"[{req_id}] Timeout")
        except Exception as e:
            logger.debug(f"[{req_id}] Error: {e}")
        finally:
            # Clean close — don't crash if already closed
            for w in (cw, uw):
                if w:
                    try:
                        w.close()
                        await w.wait_closed()
                    except Exception:
                        pass


class ProxyManager:
    """Manages auth proxy lifecycle."""
    
    _instance = None
    _proxies = {}
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    async def get_proxy_port(self, api_key: str) -> int:
        if api_key not in self._proxies:
            proxy = AuthProxy(api_key)
            await proxy.start()
            self._proxies[api_key] = proxy
        return self._proxies[api_key].port
    
    async def cleanup_proxy(self, api_key: str):
        if api_key in self._proxies:
            await self._proxies[api_key].stop()
            del self._proxies[api_key]
    
    async def cleanup_all(self):
        for proxy in list(self._proxies.values()):
            await proxy.stop()
        self._proxies.clear()
