#!/usr/bin/env python3
"""
PRE tinygrad GGUF Server — Generic OpenAI-compatible inference server.

Uses tinygrad/llm/model.py Transformer.from_gguf() for universal GGUF loading.
Supported architectures: llama, qwen2, qwen3, qwen3moe, qwen35, qwen35moe,
deepseek (MLA), glm4moe — anything tinygrad's from_gguf() handles.

Serves OpenAI-compatible /v1/chat/completions (SSE streaming) and /v1/models.

Usage:
  DEV=NV python3 tinygrad-server.py \
    --model /path/to/model.gguf \
    --tokenizer Qwen/Qwen3-32B \
    --port 8091 \
    --max-context 8192
"""

import argparse, json, time, random, sys, os, threading, traceback

# ── tinygrad imports ─────────────────────────────────────────────────────────
from pathlib import Path
from tinygrad import Tensor, Device, getenv
from tinygrad.llm.model import Transformer
from tinygrad.nn.state import get_parameters

# ── tokenizer ────────────────────────────────────────────────────────────────
try:
    from transformers import AutoTokenizer
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False

# ── HTTP server (bottle is bundled with tinygrad examples) ───────────────────
try:
    import bottle as _bottle
    from bottle import Bottle, request, response, HTTPResponse, abort
    _bottle.DEBUG = True  # render tracebacks in 500 response bodies
except ImportError:
    print("ERROR: bottle not found. Install via: pip3 install bottle", file=sys.stderr)
    sys.exit(1)


# ═══════════════════════════════════════════════════════════════════════════════
# Globals
# ═══════════════════════════════════════════════════════════════════════════════

app = Bottle()
model = None          # type: Transformer
tokenizer = None      # type: AutoTokenizer
kv_data = None        # GGUF key-value metadata
eos_token_ids = set()
model_name = "local"
model_lock = threading.Lock()  # one inference at a time (single GPU)

# CORS headers for local dev
_CORS = {
    "Access-Control-Allow-Origin": "*",
    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
    "Access-Control-Allow-Headers": "Content-Type, Authorization",
}

@app.hook("before_request")
def handle_options():
    if request.method == "OPTIONS":
        raise HTTPResponse(headers=_CORS)

@app.hook("after_request")
def enable_cors():
    for k, v in _CORS.items():
        response.set_header(k, v)


# ═══════════════════════════════════════════════════════════════════════════════
# Endpoints
# ═══════════════════════════════════════════════════════════════════════════════

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


@app.post("/v1/chat/completions")
def chat_completions():
    rjson = json.loads(request.body.read())
    messages = rjson.get("messages", [])
    if not messages:
        abort(400, "messages required")

    stream = rjson.get("stream", False)
    if not stream:
        abort(400, "streaming required (set stream: true)")

    response.content_type = "text/event-stream"
    response.set_header("Cache-Control", "no-cache")

    temperature = float(rjson.get("temperature", 0.7))
    max_tokens = int(rjson.get("max_tokens", 4096))
    random_id = f"chatcmpl-{random.randbytes(12).hex()}"

    # Apply chat template → token IDs
    try:
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        tokens = tokenizer.encode(text)
    except Exception as e:
        abort(400, f"Tokenization failed: {e}")

    def generate():
        generated = 0
        try:
            with model_lock:
                for tok in model.generate(list(tokens), temperature=temperature):
                    if tok in eos_token_ids:
                        break
                    generated += 1
                    if generated > max_tokens:
                        break

                    chunk_text = tokenizer.decode([tok])
                    chunk = {
                        "id": random_id,
                        "object": "chat.completion.chunk",
                        "created": int(time.time()),
                        "model": model_name,
                        "choices": [{
                            "index": 0,
                            "delta": {"role": "assistant", "content": chunk_text},
                            "finish_reason": None,
                        }]
                    }
                    yield f"data: {json.dumps(chunk)}\n\n"
        except Exception:
            tb = traceback.format_exc()
            sys.stderr.write("[generate] exception during streaming:\n" + tb)
            sys.stderr.flush()
            err_chunk = {
                "id": random_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model_name,
                "error": {"message": tb, "type": "InternalServerError"},
                "choices": [{"index": 0, "delta": {}, "finish_reason": "error"}],
            }
            yield f"data: {json.dumps(err_chunk)}\n\n"
            return

        stop_chunk = {
            "id": random_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": model_name,
            "choices": [{
                "index": 0,
                "delta": {},
                "finish_reason": "stop",
            }]
        }
        yield f"data: {json.dumps(stop_chunk)}\n\n"

    return generate()


@app.get("/health")
def health():
    return json.dumps({
        "status": "ok",
        "model": model_name,
        "device": str(Device.DEFAULT),
    })


# ═══════════════════════════════════════════════════════════════════════════════
# Tokenizer helpers
# ═══════════════════════════════════════════════════════════════════════════════

# Map GGUF architecture names to likely HuggingFace tokenizer repos.
# Used as fallback when --tokenizer is "auto".
_ARCH_TOKENIZER_HINTS = {
    "llama":     "meta-llama/Llama-3.3-70B-Instruct",
    "qwen2":     "Qwen/Qwen2.5-32B-Instruct",
    "qwen3":     "Qwen/Qwen3-32B",
    "qwen3moe":  "Qwen/Qwen3-235B-A22B",
    "qwen35":    "Qwen/Qwen3.5-32B",
    "qwen35moe": "Qwen/Qwen3.5-A3B",
}


def resolve_tokenizer(tokenizer_name, arch):
    """Load the right tokenizer. Supports HF names, local paths, or 'auto'."""
    if not HAS_TRANSFORMERS:
        print("ERROR: 'transformers' package required for tokenizer.", file=sys.stderr)
        print("Install via: pip3 install transformers", file=sys.stderr)
        sys.exit(1)

    if tokenizer_name == "auto":
        hint = _ARCH_TOKENIZER_HINTS.get(arch)
        if hint:
            print(f"[tokenizer] Auto-detected arch '{arch}' → trying {hint}")
            tokenizer_name = hint
        else:
            print(f"ERROR: Cannot auto-detect tokenizer for architecture '{arch}'.", file=sys.stderr)
            print(f"Please specify --tokenizer <HuggingFace model name>", file=sys.stderr)
            sys.exit(1)

    print(f"[tokenizer] Loading: {tokenizer_name}")
    tok = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
    print(f"[tokenizer] Loaded. Vocab size: {tok.vocab_size}")
    return tok


def get_eos_tokens(tok, kv):
    """Collect all possible EOS/stop token IDs."""
    ids = set()
    if tok.eos_token_id is not None:
        ids.add(tok.eos_token_id)
    # Some models have multiple EOS tokens
    if hasattr(tok, 'additional_special_tokens_ids'):
        for tid in tok.additional_special_tokens_ids:
            ids.add(tid)
    # GGUF may specify EOS
    for key in ('tokenizer.ggml.eos_token_id', 'tokenizer.ggml.eot_token_id'):
        if key in kv:
            ids.add(kv[key])
    # Qwen uses <|endoftext|> and <|im_end|>
    for special in ('<|endoftext|>', '<|im_end|>', '<|eot_id|>', '<|end_of_text|>'):
        tid = tok.convert_tokens_to_ids(special)
        if tid is not None and tid != tok.unk_token_id:
            ids.add(tid)
    return ids


# ═══════════════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    global model, tokenizer, kv_data, eos_token_ids, model_name

    parser = argparse.ArgumentParser(
        description="PRE tinygrad GGUF server — generic OpenAI-compatible inference",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--model", type=Path, required=True,
                        help="Path to GGUF model file")
    parser.add_argument("--tokenizer", type=str, default="auto",
                        help="HuggingFace tokenizer name/path, or 'auto' to detect from GGUF arch")
    parser.add_argument("--port", type=int, default=8091,
                        help="Server port")
    parser.add_argument("--host", type=str, default="127.0.0.1",
                        help="Bind address")
    parser.add_argument("--max-context", type=int, default=8192,
                        help="Maximum context window (tokens)")
    args = parser.parse_args()

    if not args.model.exists():
        print(f"ERROR: Model file not found: {args.model}", file=sys.stderr)
        sys.exit(1)

    # ── Load model ───────────────────────────────────────────────────────────
    print(f"[model] Loading GGUF: {args.model}")
    print(f"[model] Device: {Device.DEFAULT}")
    print(f"[model] Max context: {args.max_context}")

    load_start = time.time()
    model, kv_data = Transformer.from_gguf(str(args.model), max_context=args.max_context)
    load_elapsed = time.time() - load_start

    arch = kv_data.get('general.architecture', 'unknown')
    model_name = kv_data.get('general.name', arch)
    param_bytes = sum(x.nbytes() for x in get_parameters(model))

    print(f"[model] Loaded in {load_elapsed:.1f}s")
    print(f"[model] Architecture: {arch}")
    print(f"[model] Name: {model_name}")
    print(f"[model] Parameters: {param_bytes / 1e9:.1f} GB")

    # ── Load tokenizer ───────────────────────────────────────────────────────
    tokenizer = resolve_tokenizer(args.tokenizer, arch)
    eos_token_ids = get_eos_tokens(tokenizer, kv_data)
    print(f"[tokenizer] EOS token IDs: {eos_token_ids}")

    # ── Start server ─────────────────────────────────────────────────────────
    print(f"\n[server] Listening on http://{args.host}:{args.port}")
    print(f"[server] Endpoints: /v1/models, /v1/chat/completions, /health\n")
    app.run(host=args.host, port=args.port, quiet=True)


if __name__ == "__main__":
    main()
