"""Qwen3.5 (LiteRT-LM) を使った Streamlit チャットアプリ。

起動: streamlit run app.py
"""

import gc
import os
import time

import litert_lm as L
import streamlit as st

MODEL_ROOT = r"H:\LLM_Models\litert-lm"
CACHE_DIR = os.path.join(MODEL_ROOT, "cache")
DEFAULT_SYSTEM = "あなたは親切で有能なアシスタントです。日本語で簡潔に答えてください。"

# 表示名 -> (モデルファイル, 思考モード対応)
MODELS = {
    "Qwen3.5-2B (int8)": (
        r"Qwen3.5-2B-litert-lm\Qwen3.5-2B_int8.litertlm", False),
    "Qwen3.5-4B (int8)": (
        r"Qwen3.5-4B-litert-lm\Qwen3.5-4B_int8.litertlm", True),
    "Qwen3.5-4B (mixed int4)": (
        r"Qwen3.5-4B-litert-lm\Qwen3.5-4B_mixed_int4.litertlm", True),
}

L.set_min_log_severity(L.LogSeverity.ERROR)

st.set_page_config(page_title="Qwen3.5 Chat (LiteRT-LM)", page_icon="💬")


# ---------------------------------------------------------------- settings
available = {
    name: (os.path.join(MODEL_ROOT, rel), thinking)
    for name, (rel, thinking) in MODELS.items()
    if os.path.isfile(os.path.join(MODEL_ROOT, rel))
}
if not available:
    st.error(f"モデルファイルが見つかりません: {MODEL_ROOT}")
    st.stop()

with st.sidebar:
    st.header("設定")
    model_name = st.selectbox("モデル", list(available))
    missing = [name for name in MODELS if name not in available]
    if missing:
        st.caption("未配置: " + "、".join(missing))
    supports_thinking = available[model_name][1]
    backend_name = st.radio("バックエンド", ["CPU", "GPU"], horizontal=True)
    thinking = st.toggle(
        "思考モード (thinking)",
        disabled=not supports_thinking,
        help="Qwen3.5-4B のみ対応。回答前に推論過程を生成します。",
    ) and supports_thinking
    thinking_budget = st.slider(
        "思考トークン上限 (-1 で無制限)", -1, 4096, 1024, 64,
        disabled=not thinking,
    )
    system_prompt = st.text_area("システムプロンプト", DEFAULT_SYSTEM, height=120)
    temperature = st.slider("temperature", 0.0, 2.0, 0.7, 0.05)
    top_p = st.slider("top_p", 0.0, 1.0, 0.9, 0.05)
    top_k = st.slider("top_k", 1, 100, 40)
    max_output_tokens = st.slider("最大出力トークン", 64, 4096, 1024, 64)
    st.caption("設定の変更は「新しい会話」を押した時点で反映されます。")
    new_chat = st.button("🆕 新しい会話", use_container_width=True)


def current_config() -> dict:
    return dict(
        model=model_name,
        backend=backend_name,
        thinking=thinking,
        thinking_budget=thinking_budget if thinking else None,
        system=system_prompt,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        max_output_tokens=max_output_tokens,
    )


def close_chat() -> None:
    """現在の Conversation / Engine を解放する。"""
    for key in ("conversation", "engine"):
        obj = st.session_state.pop(key, None)
        if obj is not None:
            try:
                obj.close()
            except Exception:
                pass
    gc.collect()


def open_chat(cfg: dict) -> None:
    """Engine と Conversation を作り直す。

    モデルカードの注意書きにより、このモデルは会話ごとに新しい Engine を
    作る必要がある（同じ Engine で新しい会話を開くと前の会話の状態が残る）。
    """
    close_chat()
    backend = L.Backend.GPU() if cfg["backend"] == "GPU" else L.Backend.CPU()
    thinking_config = None
    if available[cfg["model"]][1]:
        thinking_config = L.ThinkingConfig(
            enable_thinking=cfg["thinking"],
            thinking_token_budget=cfg["thinking_budget"] or -1,
        )
    # 出力上限には思考トークンも含まれるため、思考の分を上乗せする
    max_tokens = cfg["max_output_tokens"]
    if cfg["thinking"]:
        budget = cfg["thinking_budget"]
        max_tokens += budget if budget and budget > 0 else 4096
    os.makedirs(CACHE_DIR, exist_ok=True)
    t0 = time.time()
    engine = L.Engine(
        available[cfg["model"]][0], backend=backend, cache_dir=CACHE_DIR
    )
    conversation = engine.create_conversation(
        system_message=cfg["system"] or None,
        sampler_config=L.SamplerConfig(
            temperature=cfg["temperature"], top_p=cfg["top_p"], top_k=cfg["top_k"]
        ),
        thinking_config=thinking_config,
        max_output_tokens=max_tokens,
    )
    st.session_state.engine = engine
    st.session_state.conversation = conversation
    st.session_state.active_config = cfg
    st.session_state.messages = []
    st.session_state.load_seconds = time.time() - t0


if (
    new_chat
    or "conversation" not in st.session_state
    or st.session_state.active_config["model"] not in available
):
    with st.spinner("モデルを読み込み中…（初回は数十秒かかります）"):
        try:
            open_chat(current_config())
        except Exception as e:
            st.error(f"モデルの読み込みに失敗しました: {e}")
            st.stop()

active = st.session_state.active_config
if active != current_config():
    st.sidebar.warning("未反映の設定変更があります。")


# ---------------------------------------------------------------- chat UI
st.title("💬 Qwen3.5 Chat")
st.caption(
    f"{active['model']} ・ {active['backend']}"
    + (" ・ 思考モード" if active["thinking"] else "")
    + f" ・ 読み込み {st.session_state.load_seconds:.1f}s"
    f" ・ 使用トークン {st.session_state.conversation.token_count} / 4096"
)


def render_thought(text: str, expanded: bool = False) -> None:
    with st.expander("💭 思考過程", expanded=expanded):
        st.markdown(text)


for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        if msg.get("thought"):
            render_thought(msg["thought"])
        st.markdown(msg["content"])


def chunk_parts(chunk) -> tuple[str, str]:
    """ストリームのチャンクから (本文, 思考) のテキストを取り出す。

    litert-lm 0.18 以降は Message、0.17 以前は
    {"role": ..., "content": [{"type": "text", "text": ...}],
     "channels": {...}} の dict を返す。思考は "thought" チャンネルに入る。
    """
    if hasattr(chunk, "contents"):
        text = str(chunk.contents)
        channels = chunk.channels or {}
    else:
        content = chunk.get("content", "")
        if isinstance(content, str):
            text = content
        else:
            text = "".join(
                c.get("text", "") for c in content if c.get("type") == "text"
            )
        channels = chunk.get("channels") or {}
    return text, channels.get("thought", "")


class ThinkSplitter:
    """本文に <think>…</think> がそのまま流れてきた場合に思考を切り分ける。"""

    START, END = "<think>", "</think>"

    def __init__(self, in_think: bool):
        self.in_think = in_think
        self.buf = ""

    def feed(self, text: str) -> tuple[str, str]:
        """(本文, 思考) を返す。"""
        self.buf += text
        answer = thought = ""
        while True:
            tag = self.END if self.in_think else self.START
            idx = self.buf.find(tag)
            if idx >= 0:
                head, self.buf = self.buf[:idx], self.buf[idx + len(tag):]
            else:
                # タグの途中で途切れている可能性がある分だけ保留する
                keep = next(
                    (k for k in range(len(tag) - 1, 0, -1)
                     if self.buf.endswith(tag[:k])),
                    0,
                )
                cut = len(self.buf) - keep
                head, self.buf = self.buf[:cut], self.buf[cut:]
            if self.in_think:
                thought += head
            else:
                answer += head
            if idx < 0:
                return answer, thought
            self.in_think = not self.in_think

    def flush(self) -> tuple[str, str]:
        rest, self.buf = self.buf, ""
        return ("", rest) if self.in_think else (rest, "")


def generate(prompt: str, use_thinking: bool) -> tuple[str, str]:
    """返答をストリーミング表示し、(本文, 思考) を返す。"""
    answer = thought = ""
    thought_box = None
    if use_thinking:
        thought_box = st.expander("💭 思考過程", expanded=True).empty()
    answer_box = st.empty()
    splitter = ThinkSplitter(in_think=False)
    for chunk in st.session_state.conversation.send_message_async(prompt):
        text, channel_thought = chunk_parts(chunk)
        a, t = splitter.feed(text)
        answer += a
        thought += channel_thought + t
        if thought and thought_box is None:
            thought_box = st.expander("💭 思考過程", expanded=True).empty()
        if thought_box is not None and (channel_thought or t):
            thought_box.markdown(thought + ("" if answer else "▌"))
        if a:
            answer_box.markdown(answer + "▌")
    a, t = splitter.flush()
    answer += a
    thought += t
    if thought_box is not None:
        thought_box.markdown(thought)
    answer_box.markdown(answer)
    return answer.strip(), thought.strip()


if prompt := st.chat_input("メッセージを入力"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        t0 = time.time()
        thought = ""
        try:
            reply, thought = generate(prompt, active["thinking"])
            if not reply:
                st.warning(
                    "回答が空でした。出力トークン上限に達した可能性があります。"
                )
        except Exception as e:
            reply = f"⚠️ 生成中にエラーが発生しました: {e}"
            st.error(reply)
        st.caption(f"{time.time() - t0:.1f}s")
    st.session_state.messages.append(
        {"role": "assistant", "content": reply, "thought": thought}
    )
    st.rerun()