"""HTTP server for real PyTorch-backed shard nodes."""

from __future__ import annotations

import http.server
import json
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from typing import Any

from .model_backend import (
    InsufficientVRAMError,
    MissingModelDependencyError,
    Quantization,
    TorchModelShard,
    validate_quantization,
)
from .server import (
    _WIRE_VERSION,
    _compress_body,
    _decompress_body,
    _parse_shape,
    _validate_activation_body,
)


class _TorchHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr,
        handler,
        backend: TorchModelShard,
        tracker_mode: bool = False,
        tracker_url: str | None = None,
        route_timeout: float = 30.0,
    ):
        super().__init__(addr, handler)
        self.backend = backend
        self.received_activations = False
        self.forward_chunk_count = 0
        self.tracker_mode = tracker_mode
        self.tracker_url = tracker_url
        self.route_timeout = route_timeout


class _TorchHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def do_POST(self):
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        if self.path == "/forward":
            self._handle_forward()
        elif self.path == "/v1/infer":
            self._handle_infer()
        elif self.path == "/v1/chat/completions" and server.tracker_mode:
            self._handle_chat_completions()
        else:
            self.send_response(404)
            self.end_headers()

    def _handle_infer(self) -> None:
        body = self._read_json_body()
        if body is None:
            return
        messages = body.get("messages", [])
        prompt = ""
        if isinstance(messages, list) and messages:
            last = messages[-1]
            if isinstance(last, dict):
                prompt = str(last.get("content", ""))
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        try:
            payload = server.backend.encode_prompt(prompt)
            if server.backend.is_tail:
                text = server.backend.decode_tail(
                    server.backend.torch.frombuffer(
                        bytearray(payload.body),
                        dtype=server.backend.torch.bfloat16,
                    )
                    .reshape(payload.shape)
                    .to(server.backend.device)
                )
                self._send_json(200, {"text": text})
                return
            self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}})
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})

    def _handle_forward(self) -> None:
        content_type = self.headers.get("Content-Type", "")
        if content_type.startswith("application/json"):
            self._handle_prompt_forward()
            return
        self._handle_binary_forward()

    def _handle_prompt_forward(self) -> None:
        body = self._read_json_body()
        if body is None:
            return
        prompt = str(body.get("prompt", ""))
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        try:
            payload = server.backend.encode_prompt(prompt)
        except Exception as exc:
            self._send_json(400, {"error": str(exc)})
            return
        self._send_activation(payload)

    def _handle_binary_forward(self) -> None:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        try:
            shape = _parse_shape(self.headers.get("X-Meshnet-Shape"))
            dtype = self.headers.get("X-Meshnet-Dtype", "")
            session = self.headers["X-Meshnet-Session"]
            chunk_index = self.headers["X-Meshnet-Chunk-Index"]
            chunk_total = self.headers["X-Meshnet-Chunk-Total"]
            encoding = self.headers.get("X-Meshnet-Encoding")
            length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(length)
            raw_body = _decompress_body(body, encoding)
            _validate_activation_body(raw_body, shape, dtype)
            if dtype != "bfloat16":
                raise ValueError("real model backend requires bfloat16 activation input")
            chunk_index_value = int(chunk_index)
            chunk_total_value = int(chunk_total)
            if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value:
                raise ValueError("invalid chunk index/total")
        except (KeyError, ValueError, TypeError):
            self.send_response(400)
            self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
            self.end_headers()
            return

        server.forward_chunk_count += 1
        if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
            server.received_activations = True

        start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
        start_layer = int(start_layer_header) if start_layer_header else None

        try:
            result = server.backend.forward_bytes(
                raw_body,
                shape,
                self.headers.get("X-Meshnet-Attn-Mask"),
                self.headers.get("X-Meshnet-Position-Ids"),
                start_layer=start_layer,
            )
        except Exception as exc:
            self._send_json(500, {"error": str(exc)})
            return

        if isinstance(result, str):
            self._send_json(200, {"text": result})
            return

        response_body = _compress_body(result.body, encoding)
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(response_body)))
        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape))
        self.send_header("X-Meshnet-Dtype", "bfloat16")
        self.send_header("X-Meshnet-Session", session)
        self.send_header("X-Meshnet-Chunk-Index", chunk_index)
        self.send_header("X-Meshnet-Chunk-Total", chunk_total)
        if encoding:
            self.send_header("X-Meshnet-Encoding", encoding)
        if result.attention_mask_header:
            self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
        if result.position_ids_header:
            self.send_header("X-Meshnet-Position-Ids", result.position_ids_header)
        self.end_headers()
        self.wfile.write(response_body)

    def _send_activation(self, payload) -> None:
        body = payload.body
        self.send_response(200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
        self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape))
        self.send_header("X-Meshnet-Dtype", "bfloat16")
        if payload.attention_mask_header:
            self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header)
        if payload.position_ids_header:
            self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header)
        self.end_headers()
        self.wfile.write(body)

    def _read_json_body(self) -> dict | None:
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length) or b"{}")
        except (json.JSONDecodeError, ValueError):
            self._send_json(400, {"error": "invalid JSON body"})
            return None
        if not isinstance(body, dict):
            self._send_json(400, {"error": "JSON body must be an object"})
            return None
        return body

    def _send_json(self, status: int, data: dict) -> None:
        payload = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        try:
            self.wfile.write(payload)
        except BrokenPipeError:
            pass  # client disconnected before we could respond — not an error

    def _handle_chat_completions(self) -> None:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        body = self._read_json_body()
        if body is None:
            return
        messages = body.get("messages", [])
        if not isinstance(messages, list):
            messages = []
        stream = bool(body.get("stream", False))
        model_name = str(body.get("model", ""))
        max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
        temperature = float(body.get("temperature") or 1.0)
        top_p = float(body.get("top_p") or 1.0)

        # Fast path: this node owns the complete model — use HF generate() with KV cache.
        # Avoids the single-token-per-forward-pass limitation of the distributed path.
        if server.backend.is_head and server.backend.is_tail:
            try:
                if stream:
                    self._stream_openai_response(
                        server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
                        model_name,
                    )
                else:
                    text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
                    self._send_openai_response(text, model_name, False, messages)
            except Exception as exc:
                self._send_json(500, {"error": f"generation failed: {exc}"})
            return

        # Distributed path: autoregressive generation across shards.
        # We do N single-step forward passes (no cross-node KV cache), which is slow
        # but correct. Each step: head encodes current sequence → forwards through route
        # → tail returns the next token string → append → repeat.
        remaining_route = self._get_remaining_route(model_name)
        print(
            f"  [node] chat route model={model_name!r} max_tokens={max_tokens} "
            f"downstream={remaining_route}",
            flush=True,
        )
        if not remaining_route:
            self._send_openai_response(
                "error: no downstream route — check tracker connectivity",
                model_name, False, messages,
            )
            return

        backend = server.backend
        # Format with chat template so the model knows it's in assistant mode.
        try:
            if hasattr(backend.tokenizer, "apply_chat_template"):
                prompt_text: str = backend.tokenizer.apply_chat_template(
                    messages, add_generation_prompt=True, tokenize=False,
                )
            else:
                raise AttributeError("no apply_chat_template")
        except Exception:
            prompt_text = " ".join(
                str(m.get("content", ""))
                for m in messages
                if isinstance(m, dict) and m.get("role") == "user"
            )

        eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
        generated: list[str] = []
        current_text = prompt_text

        for _ in range(max_tokens):
            try:
                payload = backend.encode_prompt(current_text)
            except Exception as exc:
                print(f"  [node] distributed encode error: {exc}", flush=True)
                break
            token_str = self._run_downstream_pipeline(payload, remaining_route)
            if not token_str:
                break
            # Stop on error responses or EOS.
            if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
                break
            if eos_token and token_str == eos_token:
                break
            generated.append(token_str)
            current_text = current_text + token_str

        result_text = "".join(generated)
        self._send_openai_response(result_text, model_name, stream, messages)

    def _get_remaining_route(self, model: str) -> list[tuple[str, int]]:
        """Return downstream hops as (endpoint, start_layer) pairs.

        Fast path reads X-Meshnet-Route header injected by the tracker.
        Slow path queries the tracker's /v1/route endpoint as a fallback.
        start_layer tells each downstream node which layer to begin from,
        enabling correct execution when shard ranges overlap.
        """
        # Fast path: tracker pre-resolved the downstream route and injected it as a header.
        injected = self.headers.get("X-Meshnet-Route")
        if injected:
            try:
                route = json.loads(injected)
                if isinstance(route, list):
                    hops: list[tuple[str, int]] = []
                    for item in route:
                        if isinstance(item, dict):
                            hops.append((str(item["endpoint"]), int(item.get("start_layer", 0))))
                        elif isinstance(item, str):
                            hops.append((item, 0))  # backward-compat: plain string, no start_layer
                    print(f"  [node] using injected downstream route: {[ep for ep, _ in hops]}", flush=True)
                    return hops
            except (json.JSONDecodeError, TypeError, KeyError):
                pass

        # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        if server.tracker_url is None:
            return []
        route_model = getattr(server.backend, "model_id", None) or model
        try:
            url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
            with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
                route_resp = json.loads(r.read())
            own_port = server.server_address[1]
            nodes_info = route_resp.get("nodes", [])
            # nodes_info is ordered; find own node and compute start_layers post-hoc
            hops = []
            covered_up_to: int | None = None
            for node_info in nodes_info:
                ep = node_info.get("endpoint", "")
                if ep.rstrip("/").endswith(f":{own_port}"):
                    covered_up_to = node_info.get("shard_end")
                    continue
                if covered_up_to is None:
                    # Own node not found yet; use node's shard_start as fallback
                    covered_up_to = (node_info.get("shard_start") or 1) - 1
                start_l = covered_up_to + 1
                hops.append((ep, start_l))
                covered_up_to = node_info.get("shard_end", covered_up_to)
            print(f"  [node] tracker downstream route: {[ep for ep, _ in hops]}", flush=True)
            return hops
        except Exception as exc:
            print(f"  [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
            return []

    def _run_downstream_pipeline(self, payload: object, route: list[tuple[str, int]]) -> str:
        server: _TorchHTTPServer = self.server  # type: ignore[assignment]
        if not route:
            # Partial shard at tail: decode the activation from the previous node.
            # Full single-node (head+tail) is handled before entering this method.
            if server.backend.is_tail:
                try:
                    tensor = server.backend.torch.frombuffer(
                        bytearray(payload.body),  # type: ignore[union-attr]
                        dtype=server.backend.torch.bfloat16,
                    ).reshape(payload.shape).to(server.backend.device)  # type: ignore[union-attr]
                    return server.backend.decode_tail(tensor)
                except Exception as exc:
                    return f"decode error: {exc}"
            return "no downstream route available for non-tail shard"

        session = str(uuid.uuid4())
        shape = payload.shape  # type: ignore[union-attr]
        attn_mask = payload.attention_mask_header  # type: ignore[union-attr]
        pos_ids = payload.position_ids_header  # type: ignore[union-attr]
        current_body = payload.body  # type: ignore[union-attr]
        current_shape = shape
        current_attn = attn_mask
        current_pos = pos_ids

        for hop_index, (node_url, start_layer) in enumerate(route):
            print(f"  [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True)
            headers: dict[str, str] = {
                "Content-Type": "application/octet-stream",
                "X-Meshnet-Wire": _WIRE_VERSION,
                "X-Meshnet-Shape": ",".join(str(d) for d in current_shape),
                "X-Meshnet-Dtype": "bfloat16",
                "X-Meshnet-Session": session,
                "X-Meshnet-Chunk-Index": "0",
                "X-Meshnet-Chunk-Total": "1",
                "X-Meshnet-Hop-Index": str(hop_index),
                "X-Meshnet-Start-Layer": str(start_layer),
            }
            if current_attn:
                headers["X-Meshnet-Attn-Mask"] = current_attn
            if current_pos:
                headers["X-Meshnet-Position-Ids"] = current_pos
            req = urllib.request.Request(
                f"{node_url}/forward",
                data=current_body,
                headers=headers,
                method="POST",
            )
            try:
                with urllib.request.urlopen(req, timeout=120.0) as r:
                    resp_body = r.read()
                    resp_headers = {k.lower(): v for k, v in r.headers.items()}
            except Exception as exc:
                print(f"  [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
                return f"pipeline error at {node_url}: {exc}"
            content_type = resp_headers.get("content-type", "")
            if "application/json" in content_type:
                try:
                    data = json.loads(resp_body)
                    text = str(data.get("text", ""))
                    print(f"  [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
                    return text
                except json.JSONDecodeError:
                    return resp_body.decode("utf-8", errors="replace")
            # Binary activation — update and forward to next node
            shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
            current_shape = _parse_shape(shape_header)
            current_body = resp_body
            current_attn = resp_headers.get("x-meshnet-attn-mask")
            current_pos = resp_headers.get("x-meshnet-position-ids")
        return ""

    def _stream_openai_response(self, token_iter, model: str) -> None:
        """Stream tokens from an iterator as SSE chunks."""
        chunk_id = "chatcmpl-node"
        created = int(time.time())
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()

        def _emit(data: str) -> None:
            try:
                self.wfile.write(f"data: {data}\n\n".encode())
                self.wfile.flush()
            except BrokenPipeError:
                pass

        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))
        for token_text in token_iter:
            if not token_text:
                continue
            _emit(json.dumps({
                "id": chunk_id, "object": "chat.completion.chunk", "created": created,
                "model": model,
                "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
            }))
        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }))
        try:
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
        except BrokenPipeError:
            pass

    def _send_openai_response(
        self,
        text: str,
        model: str,
        stream: bool,
        messages: list[dict] | None = None,
    ) -> None:
        chunk_id = "chatcmpl-node"
        created = int(time.time())
        if not stream:
            usage = _usage_for_response(self.server.backend, messages or [], text)  # type: ignore[attr-defined]
            self._send_json(200, {
                "id": chunk_id,
                "object": "chat.completion",
                "created": created,
                "model": model,
                "choices": [{
                    "index": 0,
                    "message": {"role": "assistant", "content": text},
                    "finish_reason": "stop",
                }],
                "usage": usage,
            })
            return
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()

        def _emit(data: str) -> None:
            try:
                self.wfile.write(f"data: {data}\n\n".encode())
                self.wfile.flush()
            except BrokenPipeError:
                pass

        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))
        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
        }))
        _emit(json.dumps({
            "id": chunk_id, "object": "chat.completion.chunk", "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }))
        try:
            self.wfile.write(b"data: [DONE]\n\n")
            self.wfile.flush()
        except BrokenPipeError:
            pass


def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
    prompt_tokens = _backend_token_count(
        backend,
        "count_prompt_tokens",
        messages,
        fallback=_fallback_message_token_count(messages),
    )
    completion_tokens = _backend_token_count(
        backend,
        "count_text_tokens",
        completion_text,
        fallback=_fallback_text_token_count(completion_text),
    )
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
    }


def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
    method: Any = getattr(backend, method_name, None)
    if callable(method):
        try:
            return max(0, int(method(value)))
        except Exception:
            pass
    return max(0, int(fallback))


def _fallback_message_token_count(messages: list[dict]) -> int:
    text = " ".join(
        str(message.get("content", ""))
        for message in messages
        if isinstance(message, dict)
    )
    return _fallback_text_token_count(text)


def _fallback_text_token_count(text: str) -> int:
    parts = text.split()
    if parts:
        return len(parts)
    return 1 if text else 0


class TorchNodeServer:
    """HTTP server backed by a HuggingFace causal language model shard."""

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 0,
        model_id: str = "openai-community/gpt2",
        shard_start: int = 0,
        shard_end: int = 6,
        quantization: str = "bfloat16",
        backend: TorchModelShard | None = None,
        tracker_mode: bool | None = None,
        tracker_url: str | None = None,
        route_timeout: float = 30.0,
    ) -> None:
        self._host = host
        self._requested_port = port
        self._backend = backend or _load_backend(
            model_id,
            shard_start,
            shard_end,
            quantization,
        )
        # Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
        self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
        self._tracker_url = tracker_url
        self._route_timeout = route_timeout
        self._server: _TorchHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    @property
    def backend(self) -> TorchModelShard:
        return self._backend

    @property
    def received_activations(self) -> bool:
        return self._server.received_activations if self._server is not None else False

    @property
    def forward_chunk_count(self) -> int:
        return self._server.forward_chunk_count if self._server is not None else 0

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("TorchNodeServer is already running")
        self._server = _TorchHTTPServer(
            (self._host, self._requested_port),
            _TorchHandler,
            self._backend,
            self._tracker_mode,
            self._tracker_url,
            self._route_timeout,
        )
        self.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self.port

    def stop(self) -> None:
        if self._server is None:
            return
        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        self._server = None
        self._thread = None
        self.port = None


def _load_backend(
    model_id: str,
    shard_start: int,
    shard_end: int,
    quantization: str,
) -> TorchModelShard:
    from .model_backend import load_torch_shard

    quant = validate_quantization(quantization)
    try:
        return load_torch_shard(model_id, shard_start, shard_end, quant)
    except MissingModelDependencyError:
        raise
    except InsufficientVRAMError as exc:
        print(f"ERROR: {exc}", file=sys.stderr, flush=True)
        raise