import asyncio
import json
import os
import hashlib
import base64
import struct
import mimetypes

# Load Configuration
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CONFIG_PATH = os.path.join(BASE_DIR, 'config.json')

with open(CONFIG_PATH, 'r') as f:
    CONFIG = json.load(f)

HTTP_PORT = CONFIG.get('http_port', 8000)
WS_PORT = CONFIG.get('ws_port', 8001)
HOST = CONFIG.get('host', 'localhost')
FRONTEND_DIR = os.path.join(BASE_DIR, 'frontend')

# --- Helper Functions ---

def create_handshake_response(key):
    GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    hash_value = hashlib.sha1((key + GUID).encode()).digest()
    accept_key = base64.b64encode(hash_value).decode()
    return (
        "HTTP/1.1 101 Switching Protocols\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        f"Sec-WebSocket-Accept: {accept_key}\r\n"
        "\r\n"
    ).encode()

def decode_frame(data):
    if len(data) < 2: return None
    byte1 = data[0]
    byte2 = data[1]
    opcode = byte1 & 0x0F
    masked = byte2 & 0x80
    payload_len = byte2 & 0x7F
    
    if opcode == 8: return "CLOSE"
    
    head = 2
    if payload_len == 126:
        payload_len = struct.unpack(">H", data[2:4])[0]
        head += 2
    elif payload_len == 127:
        payload_len = struct.unpack(">Q", data[2:10])[0]
        head += 8
        
    if masked:
        masking_key = data[head:head+4]
        head += 4
        payload = data[head:head+payload_len]
        decoded = bytearray()
        for i in range(len(payload)):
            decoded.append(payload[i] ^ masking_key[i % 4])
        return decoded.decode('utf-8')
    else:
        return data[head:head+payload_len].decode('utf-8')

def encode_frame(message):
    payload = message.encode('utf-8')
    payload_len = len(payload)
    frame = bytearray()
    frame.append(0x81)
    if payload_len <= 125:
        frame.append(payload_len)
    elif payload_len <= 65535:
        frame.append(126)
        frame.extend(struct.pack(">H", payload_len))
    else:
        frame.append(127)
        frame.extend(struct.pack(">Q", payload_len))
    frame.extend(payload)
    return frame

# --- Async HTTP Server ---

async def handle_http_client(reader, writer):
    try:
        request_line = await reader.readline()
        if not request_line:
            return
        
        method, path, _ = request_line.decode().strip().split(' ', 2)
        headers = {}
        while True:
            line = await reader.readline()
            if line == b'\r\n': break
            key, value = line.decode().strip().split(': ', 1)
            headers[key] = value
            
        print(f"[HTTP] {method} {path}")

        # Routing
        if method == 'GET':
            if path == '/' or path == '/index.html':
                await serve_file(writer, '/index.html')
            elif path == '/script.js':
                await serve_file(writer, '/script.js')
            elif path == '/config.js':
                js_content = f"const APP_CONFIG = {json.dumps(CONFIG)};".encode()
                await send_response(writer, 200, js_content, 'application/javascript')
            elif path == '/api/hello':
                response = {"message": "Hello from AsyncIO HTTP Server!"}
                await send_response(writer, 200, json.dumps(response).encode(), 'application/json')
            else:
                await send_error(writer, 404, "Not Found")
                
        elif method == 'POST':
            if path == '/api/data':
                content_length = int(headers.get('Content-Length', 0))
                body = await reader.read(content_length)
                try:
                    data = json.loads(body.decode())
                    print(f"[HTTP] Received POST data: {data}")
                    response = {"status": "success", "received": data}
                    await send_response(writer, 200, json.dumps(response).encode(), 'application/json')
                except:
                    await send_error(writer, 400, "Invalid JSON")
            else:
                await send_error(writer, 404, "Not Found")
                
        elif method == 'OPTIONS':
            writer.write(b"HTTP/1.1 200 OK\r\n")
            writer.write(b"Access-Control-Allow-Origin: *\r\n")
            writer.write(b"Access-Control-Allow-Methods: GET, POST, OPTIONS\r\n")
            writer.write(b"Access-Control-Allow-Headers: Content-Type\r\n\r\n")
            await writer.drain()
            
    except Exception as e:
        print(f"[HTTP] Error: {e}")
    finally:
        writer.close()

async def serve_file(writer, path):
    try:
        file_path = os.path.join(FRONTEND_DIR, path.lstrip('/'))
        with open(file_path, 'rb') as f:
            content = f.read()
        mime_type, _ = mimetypes.guess_type(file_path)
        await send_response(writer, 200, content, mime_type or 'application/octet-stream')
    except FileNotFoundError:
        await send_error(writer, 404, "File Not Found")

async def send_response(writer, status_code, content, content_type):
    header = f"HTTP/1.1 {status_code} OK\r\n"
    header += f"Content-Type: {content_type}\r\n"
    header += f"Content-Length: {len(content)}\r\n"
    header += "Access-Control-Allow-Origin: *\r\n"
    header += "\r\n"
    writer.write(header.encode())
    writer.write(content)
    await writer.drain()

async def send_error(writer, status_code, message):
    content = message.encode()
    header = f"HTTP/1.1 {status_code} {message}\r\n"
    header += f"Content-Length: {len(content)}\r\n\r\n"
    writer.write(header.encode())
    writer.write(content)
    await writer.drain()

# --- Async WebSocket Server ---

async def handle_ws_client(reader, writer):
    print(f"[WS] New connection")
    try:
        # Handshake
        data = await reader.read(4096) # Read handshake request
        request = data.decode()
        headers = {}
        for line in request.split("\r\n")[1:]:
            if ": " in line:
                key, value = line.split(": ", 1)
                headers[key] = value
        
        if "Sec-WebSocket-Key" in headers:
            response = create_handshake_response(headers["Sec-WebSocket-Key"])
            writer.write(response)
            await writer.drain()
            print("[WS] Handshake successful")
            
            while True:
                # Read frame header (2 bytes min)
                header = await reader.read(2)
                if not header: break
                
                byte2 = header[1]
                payload_len = byte2 & 0x7F
                
                extra_len = 0
                if payload_len == 126: extra_len = 2
                elif payload_len == 127: extra_len = 8
                
                if extra_len:
                    extended_len = await reader.read(extra_len)
                    header += extended_len
                    
                masking_key = await reader.read(4)
                header += masking_key
                
                # Calculate total length to read for payload
                if payload_len == 126:
                    payload_len = struct.unpack(">H", header[2:4])[0]
                elif payload_len == 127:
                    payload_len = struct.unpack(">Q", header[2:10])[0]
                
                payload = await reader.read(payload_len)
                full_frame = header + payload # Reconstruct for decode_frame helper
                
                message = decode_frame(full_frame)
                if message == "CLOSE":
                    print("[WS] Client requested close")
                    break
                
                if message:
                    print(f"[WS] Received: {message}")
                    response_msg = f"Async Server received: {message}"
                    writer.write(encode_frame(response_msg))
                    await writer.drain()
        else:
            print("[WS] Invalid handshake")
            
    except Exception as e:
        print(f"[WS] Error: {e}")
    finally:
        writer.close()

async def main():
    # Start HTTP Server
    http_server = await asyncio.start_server(
        handle_http_client, '0.0.0.0', HTTP_PORT)
    print(f"[HTTP] AsyncIO Serving at http://{HOST}:{HTTP_PORT}")

    # Start WebSocket Server
    ws_server = await asyncio.start_server(
        handle_ws_client, '0.0.0.0', WS_PORT)
    print(f"[WS] AsyncIO WebSocket Server listening on port {WS_PORT}")

    async with http_server, ws_server:
        await asyncio.gather(
            http_server.serve_forever(),
            ws_server.serve_forever()
        )

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nServer stopped")