"""OpenAI-compatible HTTP server for the Future-Entropy Sampler prototype -
meant to run alongside the stock qwen3.5-llama-start server (port 30000) so
the two can be compared side by side from the same chat UI (e.g. Open WebUI
as two separate model connections).

Sampler configuration (alpha/sine, top_k, confidence_threshold, ...) is fixed
at startup via CLI flags, same spirit as entropy_cli.py - there is no
per-request override and no runtime reconfiguration; restart with different
flags to try a different configuration.

Usage:
    python entropy_server.py --port 30001
    python entropy_server.py --port 30001 --sine 16,0.6,-0.2 --confidence-threshold 0.8
"""
import argparse
import json
import sys
import threading
import time
import uuid

import gguf
import jinja2
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
import uvicorn

import llama_capi as C
from entropy_sampler import EntropySampler, sine_alpha_schedule, MODEL_PATH
from entropy_cli import parse_sine, SAFE_PEAK_ALPHA


DEFAULT_MODEL_ID = "qwen3.5-35b-a3b-entropy"


def load_chat_template(model_path):
    reader = gguf.GGUFReader(model_path)
    field = reader.fields.get("tokenizer.chat_template")
    if field is None:
        raise RuntimeError(f"{model_path} has no embedded tokenizer.chat_template")
    return bytes(field.parts[-1]).decode("utf-8")


def _raise_exception(msg):
    # The Qwen chat template calls raise_exception(...) as a Jinja global,
    # matching the helper HF's apply_chat_template injects into its env.
    raise ValueError(msg)


def build_jinja_template(template_src):
    env = jinja2.Environment()
    env.globals["raise_exception"] = _raise_exception
    env.filters["tojson"] = lambda v: json.dumps(v)
    return env.from_string(template_src)


class ChatMessage(BaseModel):
    role: str
    content: str = ""


class ChatCompletionRequest(BaseModel):
    model: str | None = None
    messages: list[ChatMessage]
    stream: bool = False
    max_tokens: int | None = None
    # Accepted for OpenAI-client compatibility, intentionally ignored: this
    # sampler's behavior is fully determined by the server's baked-in
    # alpha/sine/top_k/confidence_threshold config, not per-request params.
    temperature: float | None = None
    top_p: float | None = None


def build_arg_parser():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=30001)
    p.add_argument("--model", default=MODEL_PATH)
    p.add_argument("--model-id", default=DEFAULT_MODEL_ID,
                    help="Model name reported via /v1/models and in responses.")

    alpha_group = p.add_mutually_exclusive_group()
    alpha_group.add_argument("--alpha", type=float, default=None,
                              help="Fixed alpha in [-1, 1]. Overrides --sine if given.")
    alpha_group.add_argument("--sine", type=parse_sine, default="16,0.6,-0.2",
                              metavar="PERIOD,AMPLITUDE[,OFFSET[,PHASE]]",
                              help="Default: 16,0.6,-0.2 (peak +0.4) - the best-tested "
                                   "config this session. Keep peak (offset+amplitude) "
                                   "<= ~+0.5 for safety.")

    p.add_argument("--top-k", type=int, default=12)
    p.add_argument("--top-n-future", type=int, default=20)
    p.add_argument("--confidence-threshold", type=float, default=0.8)
    p.add_argument("--n-ctx", type=int, default=212992,
                    help="Requested total context, split n_seq_max=(top_k+1) ways since "
                         "this model's KV cache is not unified across fork sequences "
                         "(kv_unified=false). Default 212992 -> ~16384 tokens/sequence "
                         "at top_k=12. Since each chat request re-sends the full "
                         "conversation history and re-primes from scratch, this is the "
                         "budget for prompt+reply *per request*, and it must also cover "
                         "the whole growing history in a multi-turn conversation.")
    p.add_argument("--n-batch", type=int, default=16384,
                    help="Hard ceiling on tokens submitted in one llama_decode() call. "
                         "prime() decodes the whole rendered prompt/history in a single "
                         "call, so this must cover the largest expected prompt. Default "
                         "16384 matches the default per-sequence context budget. Too "
                         "small and long prompts trigger GGML_ASSERT(n_tokens_all <= "
                         "n_batch), which aborts the whole process (not a catchable "
                         "Python exception) - this happened with the llama.cpp default "
                         "of 2048 on a 2059-token prompt.")
    p.add_argument("--n-ubatch", type=int, default=512,
                    help="Internal per-step chunk size within a decode call - llama.cpp "
                         "processes n_batch in n_ubatch-sized pieces automatically. "
                         "Smaller = less peak compute-buffer memory.")
    p.add_argument("--n-threads", type=int, default=8)
    p.add_argument("--max-tokens-default", type=int, default=3000,
                    help="Used when a request doesn't specify max_tokens (Open WebUI "
                         "typically doesn't). Sized for ~1600-word creative-writing "
                         "replies (~2000-2200 tokens) with margin - 512 and 1536 both "
                         "proved too tight and truncated mid-sentence.")
    p.add_argument("--think-ban", dest="think_ban", action="store_true", default=True)
    p.add_argument("--no-think-ban", dest="think_ban", action="store_false")
    p.add_argument("--enable-thinking", dest="enable_thinking", action="store_true", default=False,
                    help="Chat template's enable_thinking flag. Default off - this "
                         "sampler is tuned for direct creative-writing continuation, "
                         "not reasoning mode.")
    return p


def make_alpha(args):
    if args.alpha is not None:
        if args.alpha >= 0.7:
            print(f"warning: fixed alpha {args.alpha:+.2f} >= +0.7 reliably degenerated "
                  f"into register-breaks in prior testing", file=sys.stderr)
        return lambda: args.alpha, f"fixed={args.alpha:+.2f}"

    s = args.sine
    peak = s["offset"] + s["amplitude"]
    if peak > SAFE_PEAK_ALPHA:
        print(f"warning: peak alpha {peak:+.2f} exceeds the ~+{SAFE_PEAK_ALPHA:.1f} "
              f"safety margin observed for coherence - expect elevated crash risk",
              file=sys.stderr)
    desc = (f"sine(period={s['period_tokens']:g}, amplitude={s['amplitude']:.2f}, "
            f"offset={s['offset']:.2f}, phase={s['phase']:.2f}) peak={peak:+.2f}")
    return (lambda: sine_alpha_schedule(**s)), desc


def main():
    args = build_arg_parser().parse_args()
    alpha_factory, alpha_desc = make_alpha(args)

    print("=== entropy server ===")
    print(f"model:           {args.model}")
    print(f"model_id:        {args.model_id}")
    print(f"top_k:           {args.top_k}  (n_seq_max={args.top_k + 1})")
    print(f"top_n_future:    {args.top_n_future}")
    print(f"alpha:           {alpha_desc}")
    print(f"confidence_thr:  {args.confidence_threshold}")
    print(f"think_ban:       {args.think_ban}")
    print(f"enable_thinking: {args.enable_thinking}  (chat template flag)")
    print(f"n_ctx:           requested={args.n_ctx}")
    print(f"n_batch:         {args.n_batch}  n_ubatch: {args.n_ubatch}")

    template = build_jinja_template(load_chat_template(args.model))

    sampler = EntropySampler(
        model_path=args.model,
        n_ctx=args.n_ctx,
        top_k=args.top_k,
        top_n_future=args.top_n_future,
        n_threads=args.n_threads,
        think_ban=args.think_ban,
        n_batch=args.n_batch,
        n_ubatch=args.n_ubatch,
    )
    n_ctx_actual = C.lib.llama_n_ctx(sampler.ctx)
    n_ctx_seq = C.lib.llama_n_ctx_seq(sampler.ctx)
    print(f"n_ctx:           actual={n_ctx_actual}  per-seq={n_ctx_seq}")
    print(f"think_ban ids:   {sampler.banned_token_ids}")
    print("=" * 40)

    # Each request re-primes seq 0 from scratch with the *entire* rendered chat
    # history, so this per-sequence budget has to cover a full growing
    # conversation, not just one reply. Undersizing this is exactly what
    # produced the mid-generation "llama_decode failed rc=1 / failed to find
    # a memory slot" crash this was added after.
    if n_ctx_seq < args.max_tokens_default + 512:
        print(
            f"warning: per-sequence budget ({n_ctx_seq}) is only "
            f"{n_ctx_seq - args.max_tokens_default} tokens above max_tokens_default "
            f"({args.max_tokens_default}) - long prompts or multi-turn history will "
            f"exhaust it. Raise --n-ctx.",
            file=sys.stderr,
        )
    if args.n_batch < n_ctx_seq:
        print(
            f"warning: --n-batch ({args.n_batch}) is smaller than the per-sequence "
            f"context budget ({n_ctx_seq}) - a prompt/history longer than --n-batch "
            f"tokens will hit GGML_ASSERT(n_tokens_all <= n_batch) and abort the whole "
            f"process, even though it would otherwise fit in context. Raise --n-batch "
            f"to at least {n_ctx_seq}.",
            file=sys.stderr,
        )

    gen_lock = threading.Lock()

    def render_prompt(messages):
        try:
            return template.render(
                messages=messages, tools=None, add_generation_prompt=True,
                enable_thinking=args.enable_thinking,
            )
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))

    app = FastAPI()

    @app.get("/v1/models")
    def list_models():
        return {
            "object": "list",
            "data": [{"id": args.model_id, "object": "model", "owned_by": "local"}],
        }

    @app.post("/v1/chat/completions")
    def chat_completions(req: ChatCompletionRequest):
        messages = [m.model_dump() for m in req.messages]
        prompt = render_prompt(messages)
        max_new_tokens = req.max_tokens or args.max_tokens_default

        completion_id = f"chatcmpl-{uuid.uuid4().hex}"
        created = int(time.time())

        n_prompt_tokens = len(C.tokenize(sampler.vocab, prompt))
        print(
            f"[{completion_id}] request: prompt_tokens={n_prompt_tokens} "
            f"requested_max_tokens={req.max_tokens} resolved_max_new_tokens={max_new_tokens} "
            f"stream={req.stream}"
        )

        if req.stream:
            def event_stream():
                first = {
                    "id": completion_id, "object": "chat.completion.chunk", "created": created,
                    "model": args.model_id,
                    "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
                }
                yield f"data: {json.dumps(first)}\n\n"

                n_yielded = 0
                out_of_context = False
                with gen_lock:
                    try:
                        for piece, _diag in sampler.generate_stream(
                            prompt, max_new_tokens=max_new_tokens,
                            alpha=alpha_factory(), confidence_threshold=args.confidence_threshold,
                        ):
                            n_yielded += 1
                            chunk = {
                                "id": completion_id, "object": "chat.completion.chunk", "created": created,
                                "model": args.model_id,
                                "choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
                            }
                            yield f"data: {json.dumps(chunk)}\n\n"
                    except RuntimeError as e:
                        out_of_context = True
                        print(f"warning: generation truncated after {n_yielded} tokens: {e} "
                              f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr)

                finish_reason = "length" if (out_of_context or n_yielded >= max_new_tokens) else "stop"
                print(f"[{completion_id}] done: tokens_generated={n_yielded} finish_reason={finish_reason} "
                      f"out_of_context={out_of_context}")
                final = {
                    "id": completion_id, "object": "chat.completion.chunk", "created": created,
                    "model": args.model_id,
                    "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
                }
                yield f"data: {json.dumps(final)}\n\n"
                yield "data: [DONE]\n\n"

            return StreamingResponse(event_stream(), media_type="text/event-stream")

        out_of_context = False
        with gen_lock:
            pieces = []
            try:
                for piece, _diag in sampler.generate_stream(
                    prompt, max_new_tokens=max_new_tokens,
                    alpha=alpha_factory(), confidence_threshold=args.confidence_threshold,
                ):
                    pieces.append(piece)
            except RuntimeError as e:
                out_of_context = True
                print(f"warning: generation truncated after {len(pieces)} tokens: {e} "
                      f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr)
        text = "".join(pieces)
        finish_reason = "length" if (out_of_context or len(pieces) >= max_new_tokens) else "stop"
        print(f"[{completion_id}] done: tokens_generated={len(pieces)} finish_reason={finish_reason} "
              f"out_of_context={out_of_context}")

        return {
            "id": completion_id,
            "object": "chat.completion",
            "created": created,
            "model": args.model_id,
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": text},
                "finish_reason": finish_reason,
            }],
            "usage": {
                "prompt_tokens": 0, "completion_tokens": len(pieces),
                "total_tokens": len(pieces),
            },
        }

    try:
        uvicorn.run(app, host=args.host, port=args.port)
    finally:
        sampler.close()


if __name__ == "__main__":
    main()