"""
Copyright (C) 2024-2026 Mo Zhou <lumin@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
import argparse
from typing import Dict, List

import torch as th
import zmq
from rich.status import Status

from . import llm
from .defaults import Config, console


class AbstractBackend:
    def __init__(self, args):
        self.llm = llm.create_llm(args)

    def listen(self, args):
        raise NotImplementedError

    def server(self):
        raise NotImplementedError


def stat_messages(messages: List[Dict], llm):
    res = llm.tok.apply_chat_template(messages, tokenize=True, return_tensors="pt")
    if hasattr(res, "size"):
        context_size = res.size(1)
    elif hasattr(res, "keys") and "input_ids" in res:
        context_size = res["input_ids"].size(1)
    else:
        context_size = len(res)
    ret = f"num_msgs={len(messages)}, ctx_size={context_size}; "
    ret += f"latest={messages[-1]}"
    return ret


class ZMQBackend(AbstractBackend):
    def __init__(self, args):
        super().__init__(args)
        self.socket = zmq.Context().socket(zmq.REP)
        binduri = args.host + ":" + str(args.port)
        self.socket.bind(binduri)
        console.log(f"ZMQBackend> bind URI {binduri}. Ready to serve.")

    def listen(self):
        while True:
            msg = self.socket.recv_json()
            yield msg

    def server(self):
        for query in self.listen():
            temperature = None
            top_p = None
            if isinstance(query, dict):
                messages = query.get("session")
                temperature = query.get("temperature")
                top_p = query.get("top_p")
            else:
                messages = query

            console.log(f"DEBUG ZMQBackend query: {query} (type: {type(query)})")
            console.log(
                f"ZMQBackend> recv query: {stat_messages(messages, self.llm)}",
                markup=False,
            )

            try:
                chunks = []
                for chunk in self.llm.generate_stream(
                    messages, temperature=temperature, top_p=top_p
                ):
                    chunks.append(chunk)
                    self.socket.send_json({"status": "chunk", "content": chunk})
                    # Wait for client's ack
                    ack = self.socket.recv_json()

                # Generation complete
                generated_text = "".join(chunks)
                messages.append({"role": "assistant", "content": generated_text})
                console.log(
                    f"ZMQBackend> send reply: {stat_messages(messages, self.llm)}",
                    markup=False,
                )
                self.socket.send_json({"status": "done", "session": messages})
            except Exception as e:
                console.print_exception()
                self.socket.send_json({"status": "error", "content": str(e)})


def create_backend(args):
    if args.backend_impl == "zmq":
        backend = ZMQBackend(args)
    else:
        raise NotImplementedError(args.backend_impl)
    return backend


if __name__ == "__main__":
    ag = argparse.ArgumentParser()
    ag.add_argument(
        "--port", "-p", type=int, default=11177, help='"11177" looks like "LLM"'
    )
    ag.add_argument("--host", type=str, default="tcp://*")
    ag.add_argument("--backend_impl", type=str, default="zmq", choices=("zmq",))
    ag.add_argument("--max_new_tokens", type=int, default=512)
    conf = Config()
    ag.add_argument("--temperature", type=float, default=conf.temperature)
    ag.add_argument("--top_p", type=float, default=conf.top_p)
    ag.add_argument("--model_id", type=str, default="Qwen/Qwen3.5-0.8B")
    ag.add_argument(
        "--device", type=str, default="cuda" if th.cuda.is_available() else "cpu"
    )
    ag.add_argument(
        "--precision", type=str, default="fp16" if th.cuda.is_available() else "4bit"
    )
    ag = ag.parse_args()
    console.log(ag)

    backend = create_backend(ag)
    try:
        backend.server()
    except KeyboardInterrupt:
        pass
    console.log("Server shut down.")
