import signal
import torch
from threading import Thread, Event
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    StoppingCriteria,
    StoppingCriteriaList,
    TextIteratorStreamer,
)
from pathlib import Path

MODEL_PATH = Path(__file__).resolve().parent
DEFAULT_SYSTEM = ""


class StopOnEvent(StoppingCriteria):
    def __init__(self, event):
        self.event = event

    def __call__(self, input_ids, scores, **kwargs):
        return self.event.is_set()


def load_model():
    print(f"Loading model from: {MODEL_PATH}")

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
        trust_remote_code=True,
    )

    model.eval()

    gen_config = GenerationConfig.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True,
    )

    print(f"Loaded generation config:\n{gen_config}")

    return tokenizer, model, gen_config


def build_messages(history, user_input, system=DEFAULT_SYSTEM):
    messages = []

    if system:
        messages.append({"role": "system", "content": system})

    for u, a in history:
        messages.append({"role": "user", "content": u})
        messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": user_input})

    return messages


def join_thread(thread, stop_event):
    stop_event.set()

    # 不再 break: 确保 daemon 线程真正结束, 避免 shutdown 阶段残留线程
    while thread.is_alive():
        try:
            thread.join(timeout=0.1)
        except KeyboardInterrupt:
            stop_event.set()


@torch.inference_mode()
def stream_chat(tokenizer, model, gen_config, history, user_input):
    messages = build_messages(history, user_input)

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    stop_event = Event()
    stopping_criteria = StoppingCriteriaList([StopOnEvent(stop_event)])

    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        generation_config=gen_config,
        stopping_criteria=stopping_criteria,
    )

    thread = Thread(
        target=model.generate,
        kwargs=generate_kwargs,
        daemon=True,
    )

    thread.start()

    response = ""
    interrupted = False

    try:
        for new_text in streamer:
            print(new_text, end="", flush=True)
            response += new_text
    except KeyboardInterrupt:
        interrupted = True
        stop_event.set()
        print("\n[Interrupted. Type exit to quit or continue chatting.]")
    finally:
        join_thread(thread, stop_event)

        if not interrupted:
            print()

    return response, interrupted


def main():
    try:
        tokenizer, model, gen_config = load_model()
    except KeyboardInterrupt:
        print("\nBye.")
        return

    print("\n=== Qwen3-Sex Chat ===")
    print("Type 'exit' / 'quit' to quit, or type 'clear' to clear history.")

    history = []

    while True:
        try:
            user_input = input("User: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            break

        if not user_input:
            continue

        if user_input.lower() in {"exit", "quit"}:
            print("Bye.")
            break

        if user_input.lower() == "clear":
            history = []
            print("[History cleared]\n")
            continue

        print("Assistant: ", end="", flush=True)

        response, interrupted = stream_chat(
            tokenizer,
            model,
            gen_config,
            history,
            user_input,
        )

        if response.strip() and not interrupted:
            history.append((user_input, response))


if __name__ == "__main__":
    try:
        main()
    finally:
        while True:
            try:
                signal.signal(signal.SIGINT, signal.SIG_IGN)
                break
            except KeyboardInterrupt:
                continue
            except (ValueError, OSError):
                break