import asyncio
import re
import uuid
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
import gradio as gr

# モジュールレベルで一度だけイベントループを作成
_loop = asyncio.new_event_loop()
asyncio.set_event_loop(_loop)

# 非同期エンジンの初期化
en_args = AsyncEngineArgs(
    model="EQUES/JPharmatron-7B-chat",
    enforce_eager=True
)
model = AsyncLLMEngine.from_engine_args(en_args)

# 非同期でトークン生成をストリーミングするジェネレータ
# 生成を止めたいラベル（全角コロンも含める）
STOP_STRINGS = [
    "\nUser:", "\nユーザ：", "\nユーザー：",
    "\nAssistant:", "\nアシスタント：",
    "\nHuman:", "\nHuman："
]
# 念のための後処理（行頭の User/ユーザ/Assistant ラベル以降を削る）
STOP_RE = re.compile(r"(?:^|\n)(?:User|ユーザ|ユーザー|Assistant|アシスタント)[:：].*", re.MULTILINE)

async def astream_generate(prompt: str):
    previous_text = ""
    params = SamplingParams(
        temperature=0.0,
        max_tokens=4096,
        stop=STOP_STRINGS,     # ★ここがポイント
    )

    async for out in model.generate(prompt, params, request_id=str(uuid.uuid4())):
        full_text = out.outputs[0].text

        # 保険：もし stop が漏れても、ラベルが出たところで切る
        m = STOP_RE.search(full_text)
        if m:
            cut = m.start()
            # 直近差分だけ出す
            chunk = full_text[len(previous_text):cut]
            if chunk:
                yield chunk
            break

        # 通常は差分だけストリーム
        chunk = full_text[len(previous_text):]
        previous_text = full_text
        if chunk:
            yield chunk

# Gradio 用の応答関数（同期ジェネレータ）
def respond(user_input, history):
    history = history or []
    # システムプロンプトと過去履歴を組み立て
    base_prompt = "以下は親切で何でも答えてくれるAIアシスタントとの会話です。\n"
    for u, b in history:
        base_prompt += f"ユーザー: {u}\nアシスタント: {b}\n"
    base_prompt += f"ユーザー: {user_input}\nアシスタント:"

    # ユーザー発話と空応答を履歴に追加
    history.append((user_input, ""))

    # 単一のイベントループを使ってストリーミング
    agen = astream_generate(base_prompt)
    try:
        while True:
            chunk = _loop.run_until_complete(agen.__anext__())
            # 最新の履歴にトークンを追加
            history[-1] = (user_input, history[-1][1] + chunk)
            yield history, history
    except StopAsyncIteration:
        return

# Gradio UI 定義
with gr.Blocks() as demo:
    gr.Markdown("# 💊 製薬に関する質問をしてみてください。")
    gr.Markdown("※ ストリーミングデモ用コードです。")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="あなたのメッセージ　（Enterで送信）")
    send = gr.Button("チャットを送信")
    clear = gr.Button("チャット履歴をクリア")
    state = gr.State([])

    # ▼ プリセット（クリックでテキストボックスに反映）
    gr.Markdown("### 🔧 すぐ使えるプリセット")
    preset_list = [
        "グレープフルーツと薬を一緒に飲んじゃだめなんですか？",
        "新薬の臨床試験（Phase I〜III）の概要を、具体例つきで簡単に教えて。",
        "ジェネリック医薬品が承認されるまでの流れを、タイムラインで解説して。"
    ]
    gr.Examples(
        examples=preset_list,
        inputs=msg,
        label="プリセット（クリックで入力欄に反映）"
    )

    # 既存の送信/クリア動作
    msg.submit(respond, [msg, state], [chatbot, state]).then(lambda: "", None, msg)
    send.click(respond, [msg, state], [chatbot, state]).then(lambda: "", None, msg)
    clear.click(lambda: ([], []), None, [chatbot, state])

# エントリポイント
def main():
    demo.launch()

if __name__ == "__main__":
    main()