#!/usr/bin/env python3
"""
Minimal HTTP ↔ CoAP proxy using aiohttp + a handcrafted UDP CoAP client.

Expose REST-ish endpoints that translate browser-friendly HTTP requests to
CoAP messages for the ESP32 server.

Install dependencies:
    python3 -m pip install aiohttp

Run (defaults shown):
    HTTP_PORT=8080 COAP_HOST=192.168.4.1 COAP_PORT=5683 python3 tools/http_coap_proxy.py

Available endpoints:
  - GET /api/coap        -> query params: host, port, method, path
  - POST /api/coap       -> JSON body: {method, path, payload?, host?, port?}
  - GET /api/discover    -> queries /.well-known/core
  - GET /healthz         -> simple liveness probe
"""

import asyncio
import json
import os
import random
import signal
import socket
import struct
import time
from typing import Dict, List, Optional, Tuple

from aiohttp import web

HTTP_HOST = os.environ.get("HTTP_HOST", "0.0.0.0")
HTTP_PORT = int(os.environ.get("HTTP_PORT", "8080"))
DEFAULT_COAP_HOST = os.environ.get("COAP_HOST", "10.41.35.27")
DEFAULT_COAP_PORT = int(os.environ.get("COAP_PORT", "5683"))
REQUEST_TIMEOUT = float(os.environ.get("PROXY_TIMEOUT", "9.0"))
COAP_VERSION = 1

METHOD_CODES: Dict[str, int] = {
    "GET": 0x01,
    "POST": 0x02,
    "PUT": 0x03,
    "DELETE": 0x04,
    "FETCH": 0x05,
    "PATCH": 0x06,
    "IPATCH": 0x07,
}

OPTION_NAMES: Dict[int, str] = {
    3: "Uri-Host",
    6: "Observe",
    7: "Uri-Port",
    8: "Location-Path",
    11: "Uri-Path",
    12: "Content-Format",
    14: "Max-Age",
    15: "Uri-Query",
    17: "Accept",
    20: "Location-Query",
}


def _resolve_target(request: web.Request) -> Dict[str, str]:
    host = (request.query.get("host") or request.match_info.get("host") or DEFAULT_COAP_HOST).strip()
    port_raw = request.query.get("port") or request.match_info.get("port") or str(DEFAULT_COAP_PORT)
    if request.can_read_body and request.content_type == "application/json":
        try:
            body = request["json_body"]
        except KeyError:
            body = {}
        if isinstance(body, dict):
            host = body.get("host", host).strip()
            port_raw = str(body.get("port", port_raw))
    if not host:
        raise web.HTTPBadRequest(text="Missing CoAP host")
    try:
        port = int(port_raw)
    except ValueError:
        raise web.HTTPBadRequest(text="Invalid CoAP port")
    if port <= 0 or port > 65535:
        raise web.HTTPBadRequest(text="Invalid CoAP port range")
    return {"host": host, "port": port}


def _normalize_path(path: str) -> str:
    if not path:
        return "/"
    return path if path.startswith("/") else "/" + path


async def _coap_request(method: str, host: str, port: int, path: str, payload: Optional[str]) -> Dict[str, str]:
    loop = asyncio.get_running_loop()
    return await loop.run_in_executor(
        None,
        lambda: _coap_request_sync(method, host, port, path, payload),
    )


def _encode_extended(value: int) -> Tuple[int, bytes]:
    if value < 13:
        return value, b""
    if value < 269:
        return 13, bytes([value - 13])
    if value < 65805:
        return 14, struct.pack("!H", value - 269)
    raise ValueError("Option delta/length too large")


def _decode_extended(nibble: int, data: bytes, index: int) -> Tuple[int, int]:
    if nibble < 13:
        return nibble, index
    if nibble == 13:
        return 13 + data[index], index + 1
    if nibble == 14:
        return 269 + struct.unpack("!H", data[index:index + 2])[0], index + 2
    raise ValueError("Unsupported extended field (value 15)")


def _build_options(path: str, payload: Optional[str], method_code: int) -> bytes:
    options: List[Tuple[int, bytes]] = []
    clean_path, _, query = path.partition('?')
    segments = [seg for seg in clean_path.split('/') if seg]
    for segment in segments:
        options.append((11, segment.encode("utf-8")))
    if query:
        for item in query.split('&'):
            if item:
                options.append((15, item.encode("utf-8")))
    if payload and method_code in (0x02, 0x03, 0x05, 0x06, 0x07):
        options.append((12, b"\x00"))  # Content-Format: text/plain

    options.sort(key=lambda item: item[0])
    encoded = bytearray()
    last_number = 0
    for number, value in options:
        delta = number - last_number
        last_number = number
        delta_encoded, delta_bytes = _encode_extended(delta)
        length_encoded, length_bytes = _encode_extended(len(value))
        encoded.append((delta_encoded << 4) | length_encoded)
        encoded.extend(delta_bytes)
        encoded.extend(length_bytes)
        encoded.extend(value)
    return bytes(encoded)


def _parse_options_and_payload(data: bytes, start: int) -> Tuple[List[Dict[str, str]], str]:
    options: List[Dict[str, str]] = []
    current_opt = 0
    i = start
    length = len(data)
    while i < length:
        if data[i] == 0xFF:
            payload_bytes = data[i + 1 :]
            return options, payload_bytes.decode("utf-8", errors="replace")
        byte = data[i]
        i += 1
        delta_raw = (byte >> 4) & 0x0F
        length_raw = byte & 0x0F
        delta, i = _decode_extended(delta_raw, data, i)
        opt_length, i = _decode_extended(length_raw, data, i)
        current_opt += delta
        value = data[i : i + opt_length]
        i += opt_length

        entry: Dict[str, str] = {"number": current_opt}
        name = OPTION_NAMES.get(current_opt)
        if name:
            entry["name"] = name
        try:
            entry["value"] = value.decode("utf-8")
        except UnicodeDecodeError:
            entry["value_hex"] = value.hex()
        options.append(entry)

    return options, ""


def _format_coap_code(code: int) -> str:
    cls = (code >> 5) & 0x07
    detail = code & 0x1F
    return f"{cls}.{detail:02d}"


def _coap_request_sync(method: str, host: str, port: int, path: str, payload: Optional[str]) -> Dict[str, str]:
    method_upper = method.upper()
    method_code = METHOD_CODES.get(method_upper)
    if method_code is None:
        raise web.HTTPBadRequest(text=f"Unsupported CoAP method {method}")

    message_id = random.randint(0, 0xFFFF)
    token = os.urandom(4)
    options = _build_options(path, payload, method_code)
    payload_bytes = payload.encode("utf-8") if payload and method_code != 0x01 else b""

    token_length = len(token)
    if token_length > 8:
        raise web.HTTPBadRequest(text="Token length must be <= 8 bytes")

    header = bytearray(4)
    header[0] = (COAP_VERSION << 6) | (0 << 4) | token_length  # Confirmable
    header[1] = method_code
    header[2:4] = struct.pack("!H", message_id)

    packet = bytes(header) + token + options
    if payload_bytes:
        packet += b"\xFF" + payload_bytes

    uri = f"coap://{host}:{port}{_normalize_path(path)}"
    deadline = time.monotonic() + REQUEST_TIMEOUT
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        sock.settimeout(REQUEST_TIMEOUT)
        sock.sendto(packet, (host, port))
        print(f"[proxy] -> {method_upper} {uri} token={token.hex()} mid={message_id}", flush=True)

        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise web.HTTPGatewayTimeout(text="Timed out waiting for CoAP response")
            sock.settimeout(remaining)
            try:
                data, remote = sock.recvfrom(2048)
            except socket.timeout as exc:
                raise web.HTTPGatewayTimeout(text="Timed out waiting for CoAP response") from exc

            if len(data) < 4:
                print("[proxy] !! received too-short datagram", flush=True)
                continue

            ver = (data[0] >> 6) & 0x03
            msg_type = (data[0] >> 4) & 0x03
            token_len = data[0] & 0x0F
            resp_code = data[1]
            resp_mid = struct.unpack("!H", data[2:4])[0]
            resp_token = data[4:4 + token_len]

            debug_info = (
                f"[proxy] ?? from {remote} raw_len={len(data)} ver={ver} type={msg_type} "
                f"token_len={token_len} code_byte={resp_code} mid={resp_mid} token={resp_token.hex()}"
            )

            if ver != COAP_VERSION:
                print(debug_info + " (ignored: wrong version)", flush=True)
                continue
            if resp_mid != message_id:
                print(debug_info + f" (ignored: mid mismatch != {message_id})", flush=True)
                continue
            if resp_token != token[:token_len]:
                print(debug_info + f" (ignored: token mismatch expected {token.hex()[:token_len*2]})", flush=True)
                continue

            payload_index = 4 + token_len

            # Some lightweight CoAP stacks (eg. Arduino coap-simple) return Decimal codes (e.g. 205)
            # instead of properly encoded 5-bit class/detail. Fix that up here.
            if resp_code >= 200:
                resp_code = ((resp_code // 100) << 5) | (resp_code % 100)

            options_list, payload_text = _parse_options_and_payload(data, payload_index)
            code_string = _format_coap_code(resp_code)
            print(f"[proxy] <- {code_string} type={msg_type} token={resp_token.hex()} mid={resp_mid}", flush=True)

            # Acknowledge separate CON responses to stop retransmissions
            if msg_type == 0:  # CON
                ack = bytearray(4)
                ack[0] = (COAP_VERSION << 6) | (2 << 4)  # ACK
                ack[1] = 0
                ack[2:4] = struct.pack("!H", resp_mid)
                sock.sendto(bytes(ack), remote)

            return {
                "coap_code": code_string,
                "payload": payload_text,
                "options": options_list,
            }


async def prepare_json_body(request: web.Request):
    if request.can_read_body and request.content_type == "application/json":
        try:
            request["json_body"] = await request.json()
        except json.JSONDecodeError:
            raise web.HTTPBadRequest(text="Invalid JSON body")


async def handle_get_coap(request: web.Request) -> web.Response:
    await prepare_json_body(request)
    target = _resolve_target(request)
    method = request.query.get("method", "GET").upper()
    path = request.query.get("path", "/")
    result = await _coap_request(method, target["host"], target["port"], path, payload=None)
    return web.json_response(result)


async def handle_post_coap(request: web.Request) -> web.Response:
    await prepare_json_body(request)
    target = _resolve_target(request)

    body = request.get("json_body", {})
    if not isinstance(body, dict):
        raise web.HTTPBadRequest(text="Body must be JSON object")
    method = str(body.get("method", request.query.get("method", "PUT"))).upper()
    path = str(body.get("path", request.query.get("path", "/")))
    payload = body.get("payload", "")
    if payload is None:
        payload = ""
    if not isinstance(payload, str):
        payload = str(payload)

    result = await _coap_request(method, target["host"], target["port"], path, payload)
    return web.json_response(result)


async def handle_discover(request: web.Request) -> web.Response:
    await prepare_json_body(request)
    target = _resolve_target(request)
    result = await _coap_request("GET", target["host"], target["port"], "/.well-known/core", payload=None)
    return web.json_response(result)


async def handle_health(_: web.Request) -> web.Response:
    return web.json_response({"ok": True})


@web.middleware
async def cors_middleware(request: web.Request, handler):
    if request.method == "OPTIONS":
        resp = web.Response()
    else:
        try:
            resp = await handler(request)
        except web.HTTPException as exc:
            resp = exc
    resp.headers["Access-Control-Allow-Origin"] = request.headers.get("Origin", "*")
    resp.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS"
    resp.headers["Access-Control-Allow-Headers"] = "Content-Type"
    resp.headers["Access-Control-Allow-Credentials"] = "false"
    return resp


def create_app() -> web.Application:
    app = web.Application(middlewares=[cors_middleware])
    app.router.add_get("/api/coap", handle_get_coap)
    app.router.add_post("/api/coap", handle_post_coap)
    app.router.add_options("/api/coap", lambda _: web.Response())
    app.router.add_get("/api/discover", handle_discover)
    app.router.add_options("/api/discover", lambda _: web.Response())
    app.router.add_get("/healthz", handle_health)
    return app


def main():
    app = create_app()
    print(f"HTTP↔CoAP proxy listening on http://{HTTP_HOST}:{HTTP_PORT}")
    print(f"Default CoAP target: coap://{DEFAULT_COAP_HOST}:{DEFAULT_COAP_PORT}")
    web.run_app(app, host=HTTP_HOST, port=HTTP_PORT, print=None)


if __name__ == "__main__":
    main()