from __future__ import annotations

import argparse
import os
from pathlib import Path

import gradio as gr
from tokenizers import Tokenizer

from nayara.generate import format_prompt, load_model, stream_generate


def main() -> None:
    p = argparse.ArgumentParser(description="NAYARA local code assistant")
    p.add_argument("checkpoint")
    p.add_argument("--host", default="127.0.0.1")
    p.add_argument("--port", type=int, default=7860)
    p.add_argument("--threads", type=int, default=max(1, (os.cpu_count() or 4) // 2))
    args = p.parse_args()
    root = Path(args.checkpoint)
    tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
    model = load_model(root, int8=True, threads=args.threads)

    def answer(prompt: str, max_tokens: int, temperature: float):
        for partial in stream_generate(model, tokenizer, format_prompt(prompt), int(max_tokens), temperature):
            yield partial

    with gr.Blocks(title="NAYARA v1.0", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# NAYARA v1.0\nLocal, from-scratch code intelligence.")
        prompt = gr.Textbox(label="Ask or start coding", lines=10, placeholder="Write a fast Python LRU cache...")
        with gr.Row():
            max_tokens = gr.Slider(16, 512, value=192, step=16, label="Max new tokens")
            temperature = gr.Slider(0, 1.2, value=0.2, step=0.05, label="Temperature")
        run = gr.Button("Generate", variant="primary")
        output = gr.Code(label="NAYARA", language=None, lines=16)
        run.click(answer, [prompt, max_tokens, temperature], output)
        prompt.submit(answer, [prompt, max_tokens, temperature], output)
    demo.queue(default_concurrency_limit=1).launch(server_name=args.host, server_port=args.port)


if __name__ == "__main__":
    main()