#!/usr/bin/env python3 import argparse import json import time from threading import Lock import uvicorn from fastapi import FastAPI from fastapi.responses import HTMLResponse, StreamingResponse from pydantic import BaseModel from adapter_contract import resolve_model_directory from model_adapter import ModelAdapter HTML_PAGE = """ __PAGE_TITLE__

__PAGE_TITLE__

__PAGE_SUBTITLE__

Ready

__READY_TEXT__

__OUTPUT_LABEL__


    

  

  


"""


class GenerationRequest(BaseModel):
    prompt: str


def render_page(interaction_mode: str, runtime_name: str) -> str:
    if interaction_mode == "completion":
        values = {
            "__PAGE_TITLE__": f"{runtime_name} Text Completion",
            "__PAGE_SUBTITLE__": "Base-model text continuation on an Arm64 CPU",
            "__INPUT_LABEL__": "Text to continue",
            "__DEFAULT_INPUT__": "Title: Why Arm64 is reshaping cloud infrastructure\n\nArm64 is becoming a preferred architecture for modern cloud workloads because",
            "__BUTTON_LABEL__": "Continue",
            "__READY_TEXT__": "Ready to continue text",
            "__OUTPUT_LABEL__": "Continuation",
            "__OUTPUT_PLACEHOLDER__": "The generated continuation will appear here.",
            "__REQUIRED_MESSAGE__": "Enter text to continue",
        }
    else:
        values = {
            "__PAGE_TITLE__": f"{runtime_name} Chatbot",
            "__PAGE_SUBTITLE__": "Instruction-tuned inference on an Arm64 CPU",
            "__INPUT_LABEL__": "Prompt",
            "__DEFAULT_INPUT__": "Write a concise technical briefing for a cloud architect explaining why Arm64 CPUs are well suited to AI inference. Cover performance per watt, deployment scale, and software portability.",
            "__BUTTON_LABEL__": "Send",
            "__READY_TEXT__": "Ready for a prompt",
            "__OUTPUT_LABEL__": "Response",
            "__OUTPUT_PLACEHOLDER__": "The generated response will appear here.",
            "__REQUIRED_MESSAGE__": "Enter a prompt",
        }

    page = HTML_PAGE
    for marker, value in values.items():
        page = page.replace(marker, value)
    return page


def create_app(adapter, max_new_tokens: int, enable_thinking: bool) -> FastAPI:
    app = FastAPI()
    generation_lock = Lock()
    page = render_page(adapter.interaction_mode, adapter.runtime_name)

    @app.get("/", response_class=HTMLResponse)
    def index() -> str:
        return page

    @app.post("/generate")
    def generate(request: GenerationRequest) -> StreamingResponse:
        def response_stream():
            with generation_lock:
                start = time.perf_counter()
                first_token_time = None
                generated_tokens = 0

                for chunk in adapter.stream(
                    request.prompt,
                    max_new_tokens,
                    enable_thinking,
                ):
                    if first_token_time is None and chunk.token_count > 0:
                        first_token_time = time.perf_counter()
                    generated_tokens += chunk.token_count
                    yield json.dumps({
                        "type": "token",
                        "text": chunk.text,
                    }) + "\n"

                finish = time.perf_counter()
                decode_time = finish - first_token_time if first_token_time is not None else 0
                decoded_tokens = max(generated_tokens - 1, 0)
                yield json.dumps({
                    "type": "metrics",
                    "tokens": generated_tokens,
                    "time_to_first_token": first_token_time - start if first_token_time is not None else 0,
                    "tokens_per_second": decoded_tokens / decode_time if decode_time > 0 else 0,
                }) + "\n"

        return StreamingResponse(response_stream(), media_type="application/x-ndjson")

    return app


def main() -> None:
    parser = argparse.ArgumentParser()
    model_source = parser.add_mutually_exclusive_group(required=True)
    model_source.add_argument("--model-id", help="Downloaded Hugging Face model repository ID")
    model_source.add_argument("--model-dir", help="Downloaded model directory")
    parser.add_argument("--prompt-style", default="auto")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--enable-thinking", action="store_true", help="Enable Qwen3 thinking mode")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()

    model_dir = resolve_model_directory(args.model_id, args.model_dir)
    adapter = ModelAdapter(model_dir, args.prompt_style, args.threads)
    adapter.validate_model_directory()
    adapter.load()

    app = create_app(adapter, args.max_new_tokens, args.enable_thinking)
    uvicorn.run(app, host=args.host, port=args.port)


if __name__ == "__main__":
    main()