# -*- coding: utf-8 -*-
import socket
import threading
import time
import sys

# Configuration
A_PORT = 8080
BROWSER_PORT = 8081

# Global state
a_device_conn = None
a_device_lock = threading.Lock()
a_device_connected = threading.Event()

# Request tracking
request_queue = {}
request_queue_lock = threading.Lock()
request_counter = 0
request_counter_lock = threading.Lock()

def get_thread_id():
    """Get thread identifier for logging"""
    return threading.current_thread().getName()

def get_request_id():
    """Generate unique request ID"""
    global request_counter
    with request_counter_lock:
        request_counter += 1
        return "REQ-%06d" % request_counter

def log_info(source, message, line_num=0):
    """Log info message with thread ID"""
    thread_id = get_thread_id()
    log_line = "[%s] [INFO] [%s] [Thread:%s] %s" % (
        time.strftime("%Y-%m-%d %H:%M:%S"),
        source,
        thread_id,
        message
    )
    print(log_line)
    sys.stdout.flush()

def log_error(source, message, line_num=0):
    """Log error message with thread ID"""
    thread_id = get_thread_id()
    log_line = "[%s] [ERROR] [%s] [Thread:%s] %s" % (
        time.strftime("%Y-%m-%d %H:%M:%S"),
        source,
        thread_id,
        message
    )
    print(log_line)
    sys.stdout.flush()

def handle_a_device(conn, addr):
    """Handle device A connection - just store and hold"""
    global a_device_conn
    
    with a_device_lock:
        if a_device_conn:
            conn.send(b"ERROR: Another A device is already connected\r\n")
            conn.close()
            log_info("DeviceA", "Rejected connection from %s:%d" % addr)
            return
        
        a_device_conn = conn
    
    log_info("DeviceA", "Sending PROXY_CONNECTED to %s:%d" % addr)
    conn.send(b"PROXY_CONNECTED\r\n")
    
    try:
        conn.settimeout(30)
        log_info("DeviceA", "Waiting for READY from %s:%d" % addr)
        response = conn.recv(1024)
        log_info("DeviceA", "Received from %s:%d: %r" % (addr[0], addr[1], response))
        
        if response and b"READY" in response:
            a_device_connected.set()
            log_info("DeviceA", "Connected: %s:%d" % addr)
        else:
            with a_device_lock:
                a_device_conn = None
            conn.close()
            log_info("DeviceA", "Failed to receive READY from %s:%d" % addr)
            return
    except Exception as e:
        log_error("DeviceA", "Error receiving READY from %s:%d: %s" % (addr[0], addr[1], str(e)))
        with a_device_lock:
            a_device_conn = None
        conn.close()
        return
    
    try:
        event = threading.Event()
        event.wait()
    except:
        pass
    finally:
        with a_device_lock:
            if a_device_conn == conn:
                a_device_conn = None
                a_device_connected.clear()
        conn.close()
        log_info("DeviceA", "Disconnected")

def process_response_from_a():
    """Process responses from device A and store in request queue"""
    global a_device_conn
    while True:
        sock = None
        try:
            a_device_connected.wait()
            
            with a_device_lock:
                if not a_device_conn:
                    time.sleep(0.1)
                    continue
                
                sock = a_device_conn
            
            response = b""
            content_length = 0
            is_chunked = False
            request_id = None
            headers_found = False
            
            while True:
                chunk = sock.recv(8192)
                if not chunk:
                    if response:
                        break
                    log_info("System", "Device A connection closed in response processor")
                    with a_device_lock:
                        a_device_conn = None
                        a_device_connected.clear()
                    return
                response += chunk
                
                if b"\r\n\r\n" in response and not headers_found:
                    headers_end = response.find(b"\r\n\r\n")
                    headers = response[:headers_end].decode('utf-8', errors='replace')
                    headers_found = True
                    
                    for line in headers.split("\r\n"):
                        if line.lower().startswith("x-proxy-request-id:"):
                            request_id = line.split(":", 1)[1].strip()
                        elif line.lower().startswith("content-length:"):
                            content_length = int(line.split(":", 1)[1].strip())
                        elif line.lower().startswith("transfer-encoding:"):
                            if "chunked" in line.lower():
                                is_chunked = True
                    
                    if content_length == 0 and not is_chunked:
                        break
                    
                    if content_length > 0:
                        body_start = headers_end + 4
                        if len(response) - body_start >= content_length:
                            break
                    elif is_chunked:
                        if response.endswith(b"0\r\n\r\n"):
                            break
                elif headers_found:
                    if content_length > 0:
                        headers_end = response.find(b"\r\n\r\n")
                        body_start = headers_end + 4
                        if len(response) - body_start >= content_length:
                            break
                    elif is_chunked:
                        if response.endswith(b"0\r\n\r\n"):
                            break
            
            if request_id:
                with request_queue_lock:
                    if request_id in request_queue:
                        request_queue[request_id]['response'] = response
                        log_info("System", "Response received for request %s" % request_id)
                    else:
                        log_error("System", "No waiting browser for request %s" % request_id)
            else:
                log_error("System", "Response without request ID")
                
        except socket.timeout:
            continue
        except Exception as e:
            log_error("System", "Error in response processor: %s" % str(e))
            time.sleep(0.1)

def handle_browser_request(browser_conn, browser_addr):
    """Handle a request from browser"""
    try:
        browser_conn.settimeout(30)
        
        data = b""
        while True:
            chunk = browser_conn.recv(8192)
            if not chunk:
                break
            data += chunk
            if b"\r\n\r\n" in data:
                headers_end = data.find(b"\r\n\r\n")
                headers = data[:headers_end].decode('utf-8', errors='replace')
                
                content_length = 0
                for line in headers.split("\r\n"):
                    if line.lower().startswith("content-length:"):
                        content_length = int(line.split(":", 1)[1].strip())
                        break
                
                body_start = headers_end + 4
                if len(data) - body_start >= content_length:
                    break
        
        if not data:
            return
        
        request_lines = data.decode('utf-8', errors='replace').split("\r\n")
        if request_lines:
            first_line = request_lines[0]
            parts = first_line.split()
            if len(parts) >= 2:
                method = parts[0]
                path = parts[1]
                
                body_size = len(data) - data.find(b"\r\n\r\n") - 4 if b"\r\n\r\n" in data else 0
        
        request_id = get_request_id()
        log_info("Browser", "Request: %s %s - Body: %d bytes - ID: %s" % (method, path, body_size, request_id))
        
        with a_device_lock:
            if not a_device_conn:
                error_response = b"HTTP/1.1 503 Service Unavailable\r\nContent-Type: text/plain\r\n\r\nDevice A not connected\r\n"
                browser_conn.sendall(error_response)
                log_error("Browser", "Device A not connected for request %s" % request_id)
                return
        
        with request_queue_lock:
            request_queue[request_id] = {
                'browser_conn': browser_conn,
                'timestamp': time.time()
            }
        
        if b"\r\n\r\n" in data:
            headers_end = data.find(b"\r\n\r\n")
            headers = data[:headers_end]
            body = data[headers_end:]
            modified_data = headers + ("\r\nX-Proxy-Request-ID: " + request_id).encode('utf-8') + body
        else:
            modified_data = data + ("\r\nX-Proxy-Request-ID: " + request_id + "\r\n\r\n").encode('utf-8')
        
        with a_device_lock:
            try:
                a_device_conn.sendall(modified_data)
                log_info("Browser", "Request %s forwarded to Device A" % request_id)
            except Exception as e:
                log_error("Browser", "Failed to send to Device A for request %s: %s" % (request_id, str(e)))
                with request_queue_lock:
                    if request_id in request_queue:
                        del request_queue[request_id]
                error_response = b"HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/plain\r\n\r\nFailed to forward request to Device A\r\n"
                browser_conn.sendall(error_response)
                return
        
        timeout = 30
        start_time = time.time()
        response = None
        
        while time.time() - start_time < timeout:
            with request_queue_lock:
                if request_id in request_queue and 'response' in request_queue[request_id]:
                    response = request_queue[request_id]['response']
                    del request_queue[request_id]
                    break
                elif request_id not in request_queue:
                    break
            time.sleep(0.1)
        
        if response:
            if isinstance(response, str):
                response = response.encode('utf-8')
            browser_conn.sendall(response)
            log_info("Browser", "Response sent for request %s" % request_id)
        else:
            with request_queue_lock:
                if request_id in request_queue:
                    del request_queue[request_id]
            log_error("Browser", "Response timeout for request %s" % request_id)
            error_response = b"HTTP/1.1 504 Gateway Timeout\r\nContent-Type: text/plain\r\n\r\nResponse timeout\r\n"
            browser_conn.sendall(error_response)
        
    except socket.timeout:
        log_error("Browser", "Request timeout from %s:%d" % browser_addr)
        error_response = b"HTTP/1.1 504 Gateway Timeout\r\nContent-Type: text/plain\r\n\r\nRequest timeout\r\n"
        browser_conn.sendall(error_response)
    except Exception as e:
        log_error("Browser", "Error handling request: %s" % str(e))
        with request_queue_lock:
            if request_id in request_queue:
                del request_queue[request_id]

def handle_browser_client(conn, addr):
    """Handle browser client connection"""
    log_info("Browser", "Connected from: %s:%d" % addr)
    handle_browser_request(conn, addr)
    log_info("Browser", "Disconnected: %s:%d" % addr)

def run_a_server():
    """Start device A listening port"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('0.0.0.0', A_PORT))
    server.listen(1)
    log_info("System", "Device A port started on %d" % A_PORT)
    
    while True:
        try:
            conn, addr = server.accept()
            t = threading.Thread(target=handle_a_device, args=(conn, addr), name="DeviceA")
            t.daemon = True
            t.start()
        except:
            break
    
    server.close()

def run_browser_server():
    """Start browser listening port"""
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(('0.0.0.0', BROWSER_PORT))
    server.listen(5)
    log_info("System", "Browser port started on %d" % BROWSER_PORT)
    
    while True:
        try:
            conn, addr = server.accept()
            t = threading.Thread(target=handle_browser_client, args=(conn, addr), name="Browser-%d" % threading.activeCount())
            t.daemon = True
            t.start()
        except:
            break
    
    server.close()

def main():
    print("=" * 60)
    print("           Proxy Server (Run on Computer B)")
    print("=" * 60)
    print("Device A port:", A_PORT)
    print("Browser port:", BROWSER_PORT)
    print("=" * 60)
    print("Waiting for connections...")
    print()
    sys.stdout.flush()
    
    t1 = threading.Thread(target=run_a_server)
    t1.daemon = True
    t1.start()
    
    t2 = threading.Thread(target=run_browser_server)
    t2.daemon = True
    t2.start()
    
    t3 = threading.Thread(target=process_response_from_a, name="ResponseProc")
    t3.daemon = True
    t3.start()
    
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down...")
        sys.exit(0)

if __name__ == "__main__":
    main()