'''
恋愛シミュレーションチャットアプリ
Streamlitの画面・メイン処理
'''
# pip install requirements.txt
# streamlit run app.py

from __future__ import annotations

import os
import re
import sys
from datetime import datetime
from pathlib import Path
from uuid import uuid4

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import streamlit as st  # Webアプリの画面作成
import anthropic  # Claude API 公式 SDK

ENV_PATH = Path(__file__).resolve().parent.parent / "env" / ".env"
try:
    from dotenv import load_dotenv
    load_dotenv(dotenv_path=ENV_PATH)
except ImportError:
    raise ValueError("Unable to load API key")

import memory as mem
import json
import ui  # LINE風UIのCSS/HTMLテンプレート
import session_paths as sp  # セッションごとのデータ保存先
from typing import Any

from persona import update_persona_text
from prompt_builder import build_system_blocks

#load model
CHAT_MODEL = os.environ.get("CHAT_MODEL", "claude-haiku-4-5")
MAX_OUTPUT_TOKENS = int(os.environ.get("MAX_OUTPUT_TOKENS", "256"))  # LINE 風短文用。Anthropic は max_tokens 必須
HISTORY_TURNS = 8  #LLMに渡す直近の往復数 (古い分は history_summary に圧縮済み)

# --- サービス設定 (環境変数 / Hugging Face Secrets で渡す) ----------------
# APP_PASSKEY: 設定すると、このパスキーを知っている人だけが利用できる。
# ANTHROPIC_API_KEY: 設定するとサーバ側のキーを使う (利用者はキー入力不要)。
#   未設定なら、利用者が自分のキーを画面で入力する。
APP_PASSKEY = os.environ.get("APP_PASSKEY", "")
SERVER_API_KEY = os.environ.get("ANTHROPIC_API_KEY", "")

ROOT = sp.ROOT

def load_json(path: Path, default: Any) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return default
    
# --- LLM ------------------------------------------------------------------
@st.cache_resource(show_spinner=False)
def get_llm(api_key: str) -> anthropic.Anthropic:
    return anthropic.Anthropic(api_key=api_key)


@st.cache_resource(show_spinner=False)
def get_summary_model():
    from data.Models.Text2History.TinySwallow import (
        MODEL_NAME,
        pick_device,
        pick_dtype,
    )
    from transformers import AutoModelForCausalLM, AutoTokenizer

    device = pick_device()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=pick_dtype(device),
    ).to(device)
    model.eval()

    return tokenizer, model, device

# --- アクセス制御 / APIキー -----------------------------------------------
@st.cache_resource(show_spinner=False)
def _cleanup_once() -> bool:
    """プロセス起動後に一度だけ、古いセッションデータを掃除する。"""
    sp.cleanup_old_sessions(max_age_hours=24.0)
    return True


def require_passkey() -> None:
    """APP_PASSKEY が設定されている場合、正しいパスキーの入力を求める。"""
    if not APP_PASSKEY:
        return
    if st.session_state.get("authed"):
        return
    st.title("🔒 ログイン")
    st.caption("このアプリを使うにはパスキーが必要です。")
    with st.form("passkey_form"):
        entered = st.text_input("パスキー", type="password")
        submitted = st.form_submit_button("入る")
    if submitted:
        if entered == APP_PASSKEY:
            st.session_state.authed = True
            st.rerun()
        else:
            st.error("パスキーが違います。")
    st.stop()


def resolve_api_key() -> str:
    """利用するAnthropic APIキーを決める。利用者入力 > サーバ既定。"""
    return st.session_state.get("user_api_key", "") or SERVER_API_KEY


def require_api_key() -> str:
    """利用可能なAPIキーを確保する。なければ利用者に入力してもらう。"""
    api_key = resolve_api_key()
    if api_key:
        return api_key
    st.title("🔑 APIキーの入力")
    st.caption(
        "会話には Anthropic (Claude) の API キーが必要です。"
        "キーはこのセッション中だけメモリ上に保持され、サーバには保存されません。"
    )
    with st.form("apikey_form"):
        entered = st.text_input("Anthropic API キー (sk-ant-...)", type="password")
        submitted = st.form_submit_button("開始する")
    if submitted:
        if entered.strip():
            st.session_state.user_api_key = entered.strip()
            st.rerun()
        else:
            st.error("APIキーを入力してください。")
    st.markdown(
        "キーは [Anthropic Console](https://console.anthropic.com/settings/keys) で取得できます。"
    )
    st.stop()


# --- 画面共通の初期化 -----------------------------------------------------
st.set_page_config(page_title="チャット", page_icon=":speech_balloon:")
st.markdown(ui.CSS_BLOCK, unsafe_allow_html=True)
_cleanup_once()

# 1) パスキーによるアクセス制御
require_passkey()

# 2) セッション専用のデータ保存先を有効化 (利用者ごとに会話を分離)
if "session_id" not in st.session_state:
    st.session_state.session_id = uuid4().hex
sp.use_session(st.session_state.session_id)

# 3) APIキーを確保
api_key = require_api_key()

if "initialized" not in st.session_state:
    st.session_state.in_setup = True
    st.session_state.memory = mem.init_user_memory()
    st.session_state.initialized = True
    st.session_state.messages = []

memory = st.session_state.memory
llm = get_llm(api_key)


def render_sidebar() -> None:
    """APIキーの切り替えやログアウトなどの操作。"""
    with st.sidebar:
        st.subheader("アカウント")
        if st.session_state.get("user_api_key"):
            st.caption("APIキー: あなたのキーを使用中")
            if st.button("APIキーを消す", use_container_width=True):
                st.session_state.pop("user_api_key", None)
                st.rerun()
        elif SERVER_API_KEY:
            st.caption("APIキー: サーバ既定のキーを使用中")
            with st.expander("自分のAPIキーを使う"):
                with st.form("override_key_form"):
                    k = st.text_input("Anthropic API キー", type="password")
                    if st.form_submit_button("使う") and k.strip():
                        st.session_state.user_api_key = k.strip()
                        st.rerun()
        if APP_PASSKEY and st.button("ログアウト", use_container_width=True):
            for key in ("authed", "user_api_key"):
                st.session_state.pop(key, None)
            st.rerun()


render_sidebar()

if st.session_state.pop("api_key_error", False):
    st.error("APIキーが正しくないようです。サイドバーの『アカウント』から設定し直してください。")

def _save_uploaded_text(uploaded_file) -> Path | None:
    """アップロードされた .txt をこのセッションの text_save/ に1つだけ保存する。"""
    text_save_dir = sp.text_save_dir()
    text_save_dir.mkdir(parents=True, exist_ok=True)

    for old_file in text_save_dir.glob("*.txt"):
        old_file.unlink()

    filename = Path(str(uploaded_file.name)).name
    file_path = text_save_dir / filename
    file_path.write_bytes(uploaded_file.getvalue())

    return file_path


def _get_saved_text_file() -> Path | None:
    txt_files = list(sp.text_save_dir().glob("*.txt"))
    return txt_files[0] if txt_files else None


def _self_name_from_memory(memory_dict: dict[str, Any]) -> str:
    profile = memory_dict.get("user_profile", {})
    return (
        profile.get("user_name")
        or profile.get("nickname")
        or "Rayta"
    )

#AI 返信を空行で分割し、複数バブルとして扱う
def _split_into_bubbles(text: str) -> list[str]:
    parts = [s.strip() for s in text.split("\n\n") if s.strip()]
    return parts or [text.strip()]

def generate_reply(memory_dict, conversation, llm, placeholder, persona_name) -> str:
    recent_conversation = conversation[-HISTORY_TURNS*2:]

    self_name = _self_name_from_memory(memory_dict)
    system_blocks = build_system_blocks(memory_dict, sp.text_save_dir(), self_name)
    msgs: list[dict[str, str]] = []
    for turn in recent_conversation:
        role = turn.get("role")
        text = turn.get("text") or turn.get("content") or ""
        if not text:
            continue
        if role in ("user", "assistant"):
            msgs.append({"role": role, "content": text})

    try:
        response = llm.messages.create(
            model=CHAT_MODEL,
            max_tokens=MAX_OUTPUT_TOKENS,
            temperature=0.85,
            system=system_blocks,
            messages=msgs,
        )
        reply = next(
            (block.text for block in response.content if block.type == "text"),
            "",
        ) or ""

        placeholder.markdown(
            ui.message_html("assistant", reply, persona_name, True),
            unsafe_allow_html=True,
        )
    except anthropic.AuthenticationError:
        st.session_state["api_key_error"] = True
        return "APIキーが正しくないみたい。左のメニューからキーを設定し直してね。"
    except anthropic.RateLimitError:
        return "ごめんね、今ちょっと API の利用上限に当たってるみたい。少し時間置いて。"
    except (anthropic.APIError, anthropic.APIConnectionError):
        return "ごめん、今うまく返事できなかった。もう一度だけ送ってみて。"
    return reply.strip() or "..."

#ユーザーが送信したけどまだAIが返信していない状態を処理
def _process_pending(memory, llm, placeholder, ai_persona) -> None:
    last_user = st.session_state.messages[-1]["content"]
    mem.add_message(memory, "user", last_user)  #会話履歴にユーザーメッセージを追加
    recent_conversation = memory.get("conversation_history", [])
    raw = generate_reply(memory, recent_conversation, llm,
        placeholder, ai_persona["name"])

    bubbles = _split_into_bubbles(raw)
    #ストリーミングで描画した内容を会話履歴に保存 (\n\n で複数バブルに分割)
    for b in bubbles:
        st.session_state.messages.append({"role": "assistant", "content": b})
        mem.add_message(memory, "assistant", b)

    mem.save_memory(memory)
    if len(memory.get("conversation_history", [])) >= mem.BATCH_SIZE:
        # 履歴の要約はローカルモデルを使う重い処理。失敗しても会話は続行する。
        try:
            summary_tokenizer, summary_model, summary_device = get_summary_model()
            mem.update_summary(
                memory,
                tokenizer=summary_tokenizer,
                model=summary_model,
                device=summary_device,
            )
            mem.save_memory(memory)
        except Exception:
            pass

# --- 画面 1: 初期設定 -----------------------------------------------------
def render_setup(memory):
    profile = memory["user_profile"]
    ai_persona = memory["ai_persona"]

    st.markdown(ui.header_html("人物設定"), unsafe_allow_html=True)
    st.subheader("相手の人物設定")
    ai_persona["name"] = st.text_input(
        "相手の名前", 
        ai_persona.get("name", ""), 
        key="setup_persona_name"
    )
    st.subheader("あなたについて")
    profile["user_name"] = st.text_input(
        "あなたの名前", 
        profile.get("user_name", ""), 
        key="setup_user_name"
    )
    profile["nickname"] = st.text_input(
        "どう呼ばれてる？（例: あなたの名前が「太郎」なら「タロちゃん」とか）",
        profile.get("nickname", ""),
        key="setup_nickname",
    )
    profile["goals"] = st.text_input(
        "このアプリで何を練習したい？（例: 恋愛会話の練習 / 雑談を続ける練習）",
        profile.get("goals", ""),
        key="setup_goals",
    )
    st.subheader("二人の関係性")
    st.caption(
        "自由に書いてください。**Ex) 知り合い / 友人 / 親しい友人 / 恋人 / 彼氏 / 幼馴染** "
    )
    ai_persona["persona"] = st.text_area(
        "相手との関係性 (例: 大学のゼミで知り合った彼氏。半年前から付き合ってる)",
        ai_persona.get("persona", ""), #初期値
        height=80,
        key="setup_ai_persona",
    )
    st.subheader("テキストファイルのアップロード")
    uploaded = st.file_uploader(
        ".txt ファイルを選択 ",
        type=["txt"],
        accept_multiple_files=False,
        key="setup_text_upload",
    )
    if uploaded and st.button("アップロード"):
        saved = _save_uploaded_text(uploaded)
        st.success(f"保存しました: {saved}")
    st.divider()
    col1, col2 = st.columns([3, 1])
    with col1:
        if st.button(
            "会話を始める", type="primary", use_container_width=True, key="setup_go"
        ):
            if not ai_persona["name"].strip():
                st.error("彼の名前を入力してください")
            else:
                mem.save_memory(memory)
                st.session_state.pop("persona_update_error", None)
                txt_file = _get_saved_text_file()
                if txt_file is not None:
                    self_name = _self_name_from_memory(memory)
                    try:
                        with st.spinner("LINE履歴から人物像を作成中..."):
                            updated_memory = update_persona_text(
                                sp.user_memory_path(),
                                txt_file,
                                self_name,
                            )
                        if isinstance(updated_memory, dict):
                            st.session_state.memory = mem.load_memory()
                            memory = st.session_state.memory
                    except Exception as e:
                        st.session_state.persona_update_error = (
                            f"LINE履歴の分析に失敗しました: {e}"
                        )
                    try:
                        with st.spinner("LINE履歴から会話要約を作成中..."):
                            summary_tokenizer, summary_model, summary_device = get_summary_model()
                            mem.create_summary_from_txt(
                                memory,
                                txt_file,
                                self_name,
                                max_lines=mem.BATCH_SIZE,
                                tokenizer=summary_tokenizer,
                                model=summary_model,
                                device=summary_device,
                            )
                            mem.save_memory(memory)
                    except Exception as e:
                        st.session_state.persona_update_error = (
                            f"LINE履歴からの要約生成に失敗しました: {e}"
                        )
                st.session_state.in_setup = False
                st.rerun()
    with col2:
        if st.button("リセット", use_container_width=True, key="setup_reset"):
            import shutil
            text_save_dir = sp.text_save_dir()
            shutil.rmtree(text_save_dir, ignore_errors=True)
            text_save_dir.mkdir(parents=True, exist_ok=True)
            st.session_state.memory = mem.init_user_memory()
            st.session_state.messages = []
            st.session_state.in_setup = True
            st.rerun()

# --- 画面 2: 会話 ---------------------------------------------------------
def render_chat():
    memory = st.session_state.memory
    ai_persona = memory["ai_persona"]
    if st.session_state.get("persona_update_error"):
        st.warning(st.session_state.persona_update_error)
    st.header("設定")
    if st.button("人物設定に戻る", use_container_width=True):
        st.session_state.in_setup = True
        st.rerun()

    st.divider()
    st.caption(
        f"**{ai_persona['name']}**"
    )
    st.divider()
    if st.button("会話履歴をリセット", use_container_width=True):
        st.session_state.memory = mem.init_user_memory(True)
        st.session_state.messages = []
        st.rerun()

    st.markdown(ui.header_html(ai_persona["name"]), unsafe_allow_html=True)

    st.markdown(
        ui.chat_html(st.session_state.messages, ai_persona["name"]),
        unsafe_allow_html=True,
    )
    if (
        st.session_state.messages
        and st.session_state.messages[-1]["role"] == "user"
    ):
        #placeholder: 最初は typing インジケータ、ストリーミング開始後は実バブルへ更新
        placeholder = st.empty()
        placeholder.markdown(
            ui.typing_html(ai_persona["name"]), unsafe_allow_html=True
        )
        _process_pending(memory, llm, placeholder, ai_persona)
        st.rerun()

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


# --- ルーティング --------
if st.session_state.in_setup:
    render_setup(memory)
else:
    render_chat()