#!/usr/bin/env python3
from __future__ import annotations

import argparse
import ctypes
import os
import sys
from contextlib import contextmanager
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run inference with the local LFM2-350M MNN export."
    )
    parser.add_argument(
        "prompt",
        nargs="?",
        help="User prompt. If omitted, the script reads from stdin.",
    )
    parser.add_argument(
        "--config",
        default="config.json",
        help="Path to the exported MNN config file. Defaults to config.json next to this script.",
    )
    parser.add_argument(
        "--system",
        default="",
        help="Optional system prompt inserted ahead of the user prompt.",
    )
    parser.add_argument(
        "--stream",
        action="store_true",
        help="Stream tokens to stdout while generating.",
    )
    parser.add_argument(
        "--raw-prompt",
        action="store_true",
        help="Treat the provided prompt as a fully formatted raw model prompt.",
    )
    parser.add_argument(
        "--tmp-path",
        default="tmp",
        help="Temporary directory passed to the MNN runtime.",
    )
    parser.add_argument(
        "--show-stats",
        action="store_true",
        help="Print prompt and generation stats to stderr after inference.",
    )
    return parser.parse_args()


def resolve_path(base_dir: Path, value: str) -> Path:
    path = Path(value)
    if path.is_absolute():
        return path
    return base_dir / path


def read_prompt(args: argparse.Namespace) -> str:
    if args.prompt is not None:
        return args.prompt
    if not sys.stdin.isatty():
        prompt = sys.stdin.read()
        if prompt:
            return prompt
    raise SystemExit("Provide a prompt argument or pipe prompt text on stdin.")


def build_prompt(user_prompt: str, system_prompt: str) -> str:
    parts = ["<|startoftext|>"]
    if system_prompt:
        parts.append(f"<|im_start|>system\n{system_prompt.rstrip()}\n<|im_end|>\n")
    parts.append(f"<|im_start|>user\n{user_prompt.rstrip()}\n<|im_end|>\n<|im_start|>assistant\n")
    return "".join(parts)


@contextmanager
def suppress_native_stdout(enabled: bool):
    if not enabled:
        yield
        return

    sys.stdout.flush()
    libc = ctypes.CDLL(None)
    libc.fflush(None)
    stdout_fd = sys.stdout.fileno()
    saved_stdout_fd = os.dup(stdout_fd)

    try:
        with open(os.devnull, "w", encoding="utf-8") as devnull:
            os.dup2(devnull.fileno(), stdout_fd)
            yield
    finally:
        libc.fflush(None)
        os.dup2(saved_stdout_fd, stdout_fd)
        os.close(saved_stdout_fd)


def main() -> int:
    args = parse_args()
    base_dir = Path(__file__).resolve().parent
    config_path = resolve_path(base_dir, args.config)
    tmp_path = resolve_path(base_dir, args.tmp_path)
    tmp_path.mkdir(parents=True, exist_ok=True)

    prompt = read_prompt(args)
    formatted_prompt = prompt if args.raw_prompt else build_prompt(prompt, args.system)

    with suppress_native_stdout(not args.stream):
        import MNN.llm as mnn_llm

        model = mnn_llm.create(str(config_path))
        model.set_config({"tmp_path": str(tmp_path), "use_template": False})
        model.load()

        if model.context.status != mnn_llm.LlmStatus.RUNNING:
            raise RuntimeError(f"Model failed to load correctly: {model.context.status}")

        result = model.response(formatted_prompt, args.stream)

    if not args.stream:
        sys.stdout.write(result)
        if result and not result.endswith("\n"):
            sys.stdout.write("\n")

    if args.show_stats:
        context = model.context
        print(
            (
                f"prompt_len={context.prompt_len} "
                f"gen_seq_len={context.gen_seq_len} "
                f"prefill_us={context.prefill_us} "
                f"decode_us={context.decode_us} "
                f"status={context.status}"
            ),
            file=sys.stderr,
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())