import os
import threading
import gradio as gr
from llama_cpp import Llama

# 同時リクエスト時にCPU資源が衝突し、サーバーがフリーズするのを防ぐためのロック（排他制御）
model_lock = threading.Lock()

# 2 vCPUのCPU Basic環境に最適化
n_threads = 2

# モデルのロード
# 複数スレッドからの安全なアクセスに耐えるよう初期化
llm = Llama.from_pretrained(
    repo_id="mmnga/RakutenAI-2.0-mini-instruct-gguf",
    filename="RakutenAI-2.0-mini-instruct-Q4_0.gguf",
    n_ctx=1536,           # 長文コードに対応するための容量
    n_threads=n_threads,
    n_threads_batch=n_threads,
    n_batch=128,
    verbose=False
)

def respond(message, chat_history):
    # 他の同時リクエスト（別タブのユーザーなど）と同時に推論が走ってモデルがクラッシュするのを防ぎます。
    # 順番待ち（キュー）が発生した際も、前の処理が終わり次第、次の処理がクリアな状態で動き出します。
    with model_lock:
        # システムプロンプト
        messages = [
            {
                "role": "system",
                "content": (
                    "あなたは優秀で誠実な日本人アシスタントです。質問に対して常に丁寧かつ分かりやすく回答してください。 "
                    "プログラミングのソースコードを生成・修正する際は、途中で省略せず、最後まで完成された動作可能なコードを提示してください。"
                )
            }
        ]
        
        # メモリの引き継ぎを防ぐため、履歴はクリーンに初期化し、直近の1往復（最新の1ターン）のみを安全に抽出
        last_turns = chat_history[-1:] if chat_history else []
        
        for turn in last_turns:
            if isinstance(turn, dict):
                role = turn.get("role")
                content = turn.get("content")
                if role in ["user", "assistant"] and content:
                    messages.append({"role": role, "content": content})
            elif isinstance(turn, (list, tuple)) and len(turn) == 2:
                user_msg, bot_msg = turn
                if user_msg:
                    messages.append({"role": "user", "content": user_msg})
                if bot_msg:
                    messages.append({"role": "assistant", "content": bot_msg})

        # 最新のメッセージを追加
        messages.append({"role": "user", "content": message})

        # 推論（ストリーミング出力）
        response_stream = llm.create_chat_completion(
            messages=messages,
            max_tokens=1024,
            temperature=0.2,
            stream=True
        )
        
        partial_reply = ""
        for chunk in response_stream:
            delta = chunk["choices"][0]["delta"]
            if "content" in delta:
                partial_reply += delta["content"]
                yield partial_reply

# チャットUIの定義
demo = gr.ChatInterface(
    fn=respond,
    concurrency_limit=3,  # キューイング制限を3に設定し、CPU Basic環境でのリソース枯渇を防止
    title="Meronly - Rakuten AI 2.0 Mini",
    description="""
    低速ですがサードパーティ製Rakuten AI利用可能サイト、唯一（おそらく）のページです。<br>
    ※2 vCPU環境向けに処理を大幅に高速化・安定化させています。<br><br>
    <strong>ライセンス情報:</strong><br>
    本アプリは、Apache License 2.0 に基づき提供されている楽天グループ株式会社のオープンモデル
    <a href="https://huggingface.co/Rakuten/RakutenAI-2.0-mini-instruct" target="_blank" style="text-decoration: underline;">
    RakutenAI-2.0-mini-instruct
    </a> をコミュニティ有志が軽量化したGGUF版を使用しています。
    """
)

if __name__ == "__main__":
    demo.launch()