import sys
from core import llm_room as ap
from config.llm import (
    MODEL, N_THREADS, N_CTX, N_GPU_LAYERS, MAX_TOKENS, STREAM,
    TEMPERATURE, TOP_P, TOP_K, SYSTEM_PROMPT, CHAT_TEMPLATE,
)
from llama_cpp import Llama

def load_llm():
    print("Loading model...")
    llm = Llama(
        model_path   = str(MODEL),
        n_ctx        = N_CTX,
        n_threads    = N_THREADS,
        n_gpu_layers = N_GPU_LAYERS,
        chat_format  = CHAT_TEMPLATE or None,
        verbose      = False,
    )
    print("Model ready!")
    return llm

def chat_once(llm, messages, max_tokens):
    if STREAM:
        text = ""
        for chunk in llm.create_chat_completion(
            messages,
            stream     = True,
            temperature = TEMPERATURE,
            top_p      = TOP_P,
            top_k      = TOP_K,
            max_tokens = max_tokens,
        ):
            delta = chunk["choices"][0]["delta"].get("content")
            if delta:
                print(delta, end="", flush=True)
                text += delta
        print()
        return text

    resp = llm.create_chat_completion(
        messages,
        stream     = False,
        temperature = TEMPERATURE,
        top_p      = TOP_P,
        top_k      = TOP_K,
        max_tokens = max_tokens,
    )
    return resp["choices"][0]["message"]["content"]

def llm_run(args):

    if not MODEL.is_file():
        print(f"Model not found: {MODEL}. Download it first (see README).", file=sys.stderr)
        return

    llm        = load_llm()
    system     = args.system if args.system is not None else SYSTEM_PROMPT
    max_tokens = args.max_tokens if args.max_tokens is not None else MAX_TOKENS

    if args.prompt is not None:
        messages = [{"role": "system", "content": system}] if system else []
        messages.append({"role": "user", "content": args.prompt})
        chat_once(llm, messages, max_tokens)
        return

    messages = [{"role": "system", "content": system}] if system else []
    print("Interactive chat. Type 'exit' or 'quit' to leave.")
    try:
        while True:
            user = input("\nYou      : ").strip()
            if user.lower() in ("exit", "quit", "q"):
                print("Bye.")
                break
            if not user:
                continue
            messages.append({"role": "user", "content": user})
            print("Assistant: ", end="", flush=True)
            reply = chat_once(llm, messages, max_tokens)
            messages.append({"role": "assistant", "content": reply})
    except (KeyboardInterrupt, EOFError):
        print("\nBye.")

if __name__ == "__main__":
    llm_run( ap.parser.parse_args() )