#!/usr/bin/env python3
"""
Sixpert K1 - Example Generation Script
=======================================
Demonstrates how to load and run inference with Sixpert K1 (Q4_K_M GGUF)
using llama-cpp-python.

Usage:
    pip install llama-cpp-python
    python generate.py --prompt "Explain quantum entanglement"
"""

import argparse
import time
import sys

try:
    from llama_cpp import Llama
except ImportError:
    print("Installing llama-cpp-python...")
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
    from llama_cpp import Llama


def format_prompt(messages: list[dict]) -> str:
    """Format messages into Sixpert chat template."""
    formatted = ""
    for msg in messages:
        role = msg["role"]
        content = msg["content"]
        if role == "system":
            formatted += f"<|im_start|>system\n{content}<|im_end|>\n"
        elif role == "user":
            formatted += f"<|im_start|>user\n{content}<|im_end|>\n"
        elif role == "assistant":
            formatted += f"<|im_start|>assistant\n{content}<|im_end|>\n"
    formatted += "<|im_start|>assistant\n"
    return formatted


def run_generation(
    model_path: str,
    prompt: str,
    max_tokens: int = 2048,
    temperature: float = 0.7,
    top_p: float = 0.8,
    top_k: int = 40,
    repeat_penalty: float = 1.05,
    gpu_layers: int = -1,
    threads: int = 8,
    verbose: bool = True,
):
    """Run text generation with Sixpert K1."""

    print(f"Loading Sixpert K1 from: {model_path}")
    print(f"Quantization: Q4_K_M | Layers: {gpu_layers if gpu_layers > 0 else 'All (offload)'}")
    print("-" * 60)

    llm = Llama(
        model_path=model_path,
        n_ctx=131072,
        n_gpu_layers=gpu_layers,
        n_threads=threads,
        verbose=False,
    )

    messages = [
        {
            "role": "system",
            "content": "You are Sixpert K1, a precision logic engine developed by Sixpert AI. "
            "You excel at reasoning, code generation, mathematical analysis, "
            "and structured problem-solving. You think before you respond.",
        },
        {"role": "user", "content": prompt},
    ]

    formatted_prompt = format_prompt(messages)

    if verbose:
        print(f"\nPrompt:\n{prompt}\n")
        print("Generating response...")
        print("-" * 40)

    start_time = time.time()

    stream = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repeat_penalty=repeat_penalty,
        stream=True,
    )

    full_response = ""
    for chunk in stream:
        delta = chunk["choices"][0]["delta"].get("content", "")
        if delta:
            full_response += delta
            if verbose:
                print(delta, end="", flush=True)

    elapsed = time.time() - start_time

    if verbose:
        print("\n")
        print("-" * 60)
        print(f"Generation completed in {elapsed:.2f}s")
        print(f"Output: {len(full_response.split())} words | {len(full_response)} chars")

    return full_response


def main():
    parser = argparse.ArgumentParser(description="Sixpert K1 Generation Script")
    parser.add_argument(
        "--model", type=str, default="SixpertK1.gguf", help="Path to GGUF model file"
    )
    parser.add_argument("--prompt", type=str, default="What is your name and what can you do?", help="Input prompt")
    parser.add_argument("--max-tokens", type=int, default=2048, help="Maximum tokens to generate")
    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
    parser.add_argument("--top-p", type=float, default=0.8, help="Top-p sampling")
    parser.add_argument("--top-k", type=int, default=40, help="Top-k sampling")
    parser.add_argument("--gpu-layers", type=int, default=-1, help="GPU layers to offload (-1 for all)")
    parser.add_argument("--threads", type=int, default=8, help="CPU threads")
    parser.add_argument("--verbose", action="store_true", default=True, help="Verbose output")

    args = parser.parse_args()

    run_generation(
        model_path=args.model,
        prompt=args.prompt,
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        gpu_layers=args.gpu_layers,
        threads=args.threads,
        verbose=args.verbose,
    )


if __name__ == "__main__":
    main()