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

# Configuration from config.ini
CONFIG_FILE = os.path.join(os.path.dirname(__file__), 'config.ini')

def load_config():
    config = configparser.ConfigParser()
    config.read(CONFIG_FILE)
    
    target_host = config.get('proxy', 'target_host', fallback='zentao.sunyard.com.cn')
    target_port = config.getint('proxy', 'target_port', fallback=9788)
    target_protocol = config.get('proxy', 'target_protocol', fallback='http').lower()
    b_host = config.get('proxy', 'b_host', fallback='localhost')
    b_port = config.getint('proxy', 'b_port', fallback=8080)
    
    return target_host, target_port, target_protocol, b_host, b_port

TARGET_HOST, TARGET_PORT, TARGET_PROTOCOL, B_HOST, B_PORT = load_config()

# Global lock for socket write operations
socket_write_lock = threading.Lock()

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

def log_info(source, message):
    """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):
    """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 extract_request_id(request_data):
    """Extract X-Proxy-Request-ID from request headers"""
    if "\r\n\r\n" in request_data:
        headers_end = request_data.find("\r\n\r\n")
        headers = request_data[:headers_end]
        for line in headers.split("\r\n"):
            if line.lower().startswith("x-proxy-request-id:"):
                return line.split(":", 1)[1].strip()
    return None

def modify_request_headers(request_data):
    """Modify request headers: replace Origin, Referer, Host; remove Connection headers"""
    lines = request_data.split("\r\n")
    new_lines = []
    method = ""
    path = ""
    
    if lines:
        first_line = lines[0]
        parts = first_line.split()
        if len(parts) >= 2:
            method = parts[0]
            path = parts[1]
    
    for line in lines:
        line_str = line
        
        if line_str.startswith("Origin:"):
            new_lines.append("Origin: %s://%s:%d" % (TARGET_PROTOCOL, TARGET_HOST, TARGET_PORT))
        elif line_str.startswith("Referer:"):
            new_lines.append("Referer: %s://%s:%d/" % (TARGET_PROTOCOL, TARGET_HOST, TARGET_PORT))
        elif line_str.startswith("Host:"):
            new_lines.append("Host: %s:%d" % (TARGET_HOST, TARGET_PORT))
        elif line_str.lower().startswith("connection:"):
            continue
        elif line_str.lower().startswith("keep-alive:"):
            continue
        elif line_str.lower().startswith("proxy-connection:"):
            continue
        elif line_str.lower().startswith("x-forwarded-for:"):
            continue
        elif line_str.lower().startswith("x-proxy-request-id:"):
            continue
        else:
            new_lines.append(line_str)
    
    modified_data = "\r\n".join(new_lines)
    
    if path.startswith("/"):
        new_path = "%s://%s:%d%s" % (TARGET_PROTOCOL, TARGET_HOST, TARGET_PORT, path)
        if method and modified_data:
            modified_data = modified_data.replace(path, new_path, 1)
    
    return modified_data, method, path

def send_request_to_target(host, port, protocol, request_data, request_id):
    """Send request to target server and return response"""
    target_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    target_sock.settimeout(30)
    
    try:
        target_sock.connect((host, port))
        log_info("System", "Connected to target: %s:%d" % (host, port))
        
        if protocol == 'https':
            context = ssl.create_default_context()
            target_sock = context.wrap_socket(target_sock, server_hostname=host)
            log_info("System", "SSL connection established with target")
        
        if isinstance(request_data, str):
            request_data = request_data.encode('utf-8')
        target_sock.sendall(request_data)
        
        response = b""
        content_length = 0
        is_chunked = False
        
        while True:
            chunk = target_sock.recv(8192)
            if not chunk:
                break
            response += chunk
            
            if b"\r\n\r\n" in response and content_length == 0 and not is_chunked:
                headers_end = response.find(b"\r\n\r\n")
                headers = response[:headers_end].decode('utf-8', errors='replace')
                
                for line in headers.split("\r\n"):
                    if 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:
                    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:
            if b"\r\n\r\n" in response:
                headers_end = response.find(b"\r\n\r\n")
                headers = response[:headers_end]
                body = response[headers_end:]
                response = headers + ("\r\nX-Proxy-Request-ID: " + request_id).encode('utf-8') + body
        
        return response
    
    finally:
        target_sock.close()

def parse_location(location):
    """Parse Location header to get protocol, host, port, and path"""
    protocol = TARGET_PROTOCOL
    host = TARGET_HOST
    port = TARGET_PORT
    path = "/"
    
    if location.startswith("http://") or location.startswith("https://"):
        # Full URL
        protocol_end = location.find("://")
        protocol = location[:protocol_end]
        remaining = location[protocol_end + 3:]
        path_start = remaining.find("/")
        if path_start == -1:
            host_part = remaining
            path = "/"
        else:
            host_part = remaining[:path_start]
            path = remaining[path_start:]
        
        port_pos = host_part.find(":")
        if port_pos != -1:
            host = host_part[:port_pos]
            port = int(host_part[port_pos + 1:])
        else:
            host = host_part
            port = 443 if protocol == 'https' else 80
    elif location.startswith("/"):
        # Absolute path
        path = location
    else:
        # Relative path
        path = "/" + location
    
    return protocol, host, port, path

def modify_redirect_request(original_request, new_protocol, new_host, new_port, new_path):
    """Modify request for redirect"""
    lines = original_request.split("\r\n")
    new_lines = []
    
    first_line = lines[0]
    parts = first_line.split()
    if len(parts) >= 3:
        method = parts[0]
        http_version = parts[2]
        new_first_line = "%s %s %s" % (method, new_path, http_version)
        new_lines.append(new_first_line)
    
    for line in lines[1:]:
        if not line:
            continue
        if line.lower().startswith("host:"):
            new_lines.append("Host: %s:%d" % (new_host, new_port))
        elif line.lower().startswith("origin:"):
            new_lines.append("Origin: %s://%s:%d" % (new_protocol, new_host, new_port))
        elif line.lower().startswith("referer:"):
            new_lines.append("Referer: %s://%s:%d/" % (new_protocol, new_host, new_port))
        elif line.lower().startswith("connection:"):
            continue
        elif line.lower().startswith("keep-alive:"):
            continue
        elif line.lower().startswith("proxy-connection:"):
            continue
        elif line.lower().startswith("x-forwarded-for:"):
            continue
        elif line.lower().startswith("x-proxy-request-id:"):
            continue
        else:
            new_lines.append(line)
    
    new_lines.append("")
    new_lines.append("")
    
    if "\r\n\r\n" in original_request:
        body_start = original_request.find("\r\n\r\n") + 4
        body = original_request[body_start:]
        modified_data = "\r\n".join(new_lines[:-2]) + "\r\n\r\n" + body
    else:
        modified_data = "\r\n".join(new_lines[:-2]) + "\r\n\r\n"
    
    return modified_data

def handle_request(request_data):
    """Handle a single request from device B - completely independent thread"""
    if isinstance(request_data, bytes):
        request_data_str = request_data.decode('utf-8', errors='replace')
    else:
        request_data_str = request_data
    
    request_id = extract_request_id(request_data_str)
    
    modified_data, method, path = modify_request_headers(request_data_str)
    
    body_size = 0
    if "\r\n\r\n" in modified_data:
        body_size = len(modified_data) - modified_data.find("\r\n\r\n") - 4
    
    log_info("System", "Request: %s %s://%s:%d%s - Body: %d bytes - ID: %s" % (method, TARGET_PROTOCOL, TARGET_HOST, TARGET_PORT, path, body_size, request_id))
    
    current_protocol = TARGET_PROTOCOL
    current_host = TARGET_HOST
    current_port = TARGET_PORT
    current_request = modified_data
    redirect_count = 0
    max_redirects = 5
    
    try:
        while redirect_count <= max_redirects:
            response = send_request_to_target(current_host, current_port, current_protocol, current_request, request_id)
            
            status_code = "500"
            location = None
            
            if b"\r\n\r\n" in response:
                headers_end = response.find(b"\r\n\r\n")
                headers = response[:headers_end].decode('utf-8', errors='replace')
                for line in headers.split("\r\n"):
                    if line.startswith("HTTP/"):
                        status_code = line.split()[1]
                    elif line.lower().startswith("location:"):
                        location = line.split(":", 1)[1].strip()
            
            log_info("System", "Response: %s %s - Status: %s - ID: %s" % (method, path, status_code, request_id))
            
            if status_code not in ("301", "302") or not location:
                return response
            
            redirect_count += 1
            log_info("System", "Redirect %d: %s" % (redirect_count, location))
            
            new_protocol, new_host, new_port, new_path = parse_location(location)
            current_protocol = new_protocol
            current_host = new_host
            current_port = new_port
            current_request = modify_redirect_request(request_data_str, new_protocol, new_host, new_port, new_path)
            path = new_path
        
        log_error("System", "Too many redirects (%d)" % max_redirects)
        error_response = ("HTTP/1.1 508 Loop Detected\r\nX-Proxy-Request-ID: " + (request_id if request_id else "") + "\r\nContent-Type: text/plain\r\n\r\nToo many redirects\r\n").encode('utf-8')
        return error_response
    
    except socket.error as e:
        if e.errno == 111 or e.errno == 10061:
            log_error("System", "Cannot connect to target: %s:%d" % (current_host, current_port))
            error_response = ("HTTP/1.1 502 Bad Gateway\r\nX-Proxy-Request-ID: " + (request_id if request_id else "") + "\r\nContent-Type: text/plain\r\n\r\nCannot connect to target: %s:%d\r\n" % (current_host, current_port)).encode('utf-8')
            return error_response
        else:
            log_error("System", "Socket error: %s" % str(e))
            error_response = ("HTTP/1.1 502 Bad Gateway\r\nX-Proxy-Request-ID: " + (request_id if request_id else "") + "\r\nContent-Type: text/plain\r\n\r\nSocket error\r\n").encode('utf-8')
            return error_response
    except socket.timeout:
        log_error("System", "Target connection timeout")
        error_response = ("HTTP/1.1 504 Gateway Timeout\r\nX-Proxy-Request-ID: " + (request_id if request_id else "") + "\r\nContent-Type: text/plain\r\n\r\nConnection timeout\r\n").encode('utf-8')
        return error_response
    except Exception as e:
        log_error("System", "Error handling request: %s" % str(e))
        error_response = ("HTTP/1.1 500 Internal Server Error\r\nX-Proxy-Request-ID: " + (request_id if request_id else "") + "\r\nContent-Type: text/plain\r\n\r\nInternal error\r\n").encode('utf-8')
        return error_response

def process_request(b_socket, request_data):
    """Process a request in separate thread and send response back"""
    response = handle_request(request_data)
    
    with socket_write_lock:
        try:
            log_info("System", "Sending response back to B, length: %d bytes" % len(response))
            b_socket.sendall(response)
            log_info("System", "Response sent successfully")
        except Exception as e:
            log_error("System", "Failed to send response: %s" % str(e))

def listen_for_requests(sock):
    """Listen for requests from device B"""
    sock.settimeout(300)
    
    while True:
        try:
            data = b""
            while True:
                chunk = sock.recv(8192)
                if not chunk:
                    log_info("System", "Connection closed by device B")
                    return
                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
            
            t = threading.Thread(target=process_request, args=(sock, data), name="ReqHandler-%d" % threading.activeCount())
            t.daemon = True
            t.start()
            
        except socket.timeout:
            log_info("System", "Connection idle timeout, continuing to wait")
            continue
        except Exception as e:
            log_error("System", "Connection error: %s" % str(e))
            return

def connect_to_b_device(b_host, b_port):
    """Connect to device B and handle requests"""
    failure_count = 0
    max_failures = 3
    
    while True:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)
            log_info("System", "Connecting to device B: %s:%d" % (b_host, b_port))
            sock.connect((b_host, b_port))
            
            log_info("System", "Connected to device B, waiting for PROXY_CONNECTED")
            response = sock.recv(1024)
            log_info("System", "Received from B: %r" % response)
            
            if not response or b"PROXY_CONNECTED" not in response:
                log_error("System", "Failed to receive connection confirmation")
                sock.close()
                failure_count += 1
                if failure_count >= max_failures:
                    log_error("System", "Maximum connection attempts reached (%d), exiting..." % max_failures)
                    sys.exit(1)
                time.sleep(3)
                continue
            
            log_info("System", "Sending READY to device B")
            sock.sendall(b"READY\r\n")
            log_info("System", "Successfully connected to device B: %s:%d" % (b_host, b_port))
            
            failure_count = 0
            listen_for_requests(sock)
            
        except socket.error as e:
            if e.errno == 111 or e.errno == 10061:
                log_info("System", "Device B not ready (%s:%d), retrying in 3 seconds..." % (b_host, b_port))
            else:
                log_error("System", "Connection error: %s" % str(e))
            failure_count += 1
            if failure_count >= max_failures:
                log_error("System", "Maximum connection attempts reached (%d), exiting..." % max_failures)
                sys.exit(1)
            time.sleep(3)
        except Exception as e:
            log_error("System", "Unexpected error: %s" % str(e))
            failure_count += 1
            if failure_count >= max_failures:
                log_error("System", "Maximum connection attempts reached (%d), exiting..." % max_failures)
                sys.exit(1)
            time.sleep(3)

def main():
    print("=" * 60)
    print("           Proxy Client (Run on Computer A)")
    print("=" * 60)
    print("Target website:", TARGET_PROTOCOL, "://", TARGET_HOST, ":", TARGET_PORT)
    print("Device B address:", B_HOST, ":", B_PORT)
    print("=" * 60)
    sys.stdout.flush()
    
    try:
        connect_to_b_device(B_HOST, B_PORT)
    except KeyboardInterrupt:
        print("\nShutting down...")
        sys.exit(0)

if __name__ == "__main__":
    main()