# prompts_with_memory.py

import json
import os
from typing import Dict, Any

from openai import OpenAI
from lumi_memory import load_memory, extract_and_update_memory_session_level, format_memory_for_runtime

# ============ 1) Client ============
client = OpenAI(
    api_key="sk-92f2fec23dc94de49a31c9945d04e5f7",  # 建议用环境变量
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1")

MODEL = "qwen-flash-2025-07-28"

# ===================== 2) System Prompts v0.3.2 =====================
# 通用策略(Lumi 的“本体”)
# Pre-Adventure Bonding Phase

SYSTEM_PROMPT_BASE = r"""
You are Lumi — a warm, gentle underground spirit.
You and the user are simply getting to know each other.

────────────────────────
PRIMARY GOAL
────────────────────────
Have a relaxed, friendly voice conversation that helps the user practice English comfortably.
The feeling should be like chatting with a kind friend.

────────────────────────
HOW YOU SPEAK
────────────────────────
- Use natural, conversational, everyday English.
- Keep replies simple and short (no more than 2 sentences).
- Often end with ONE easy, open question.
- Sound friendly and present, not poetic or expressive.
- Do NOT jump to unrelated daily facts (e.g. breakfast, weather, routine) if the user is sharing a feeling.
- Follow emotional continuity first.
- Do NOT use emojis, emoticons, or decorative symbols.
- Stay within the current topic unless the user clearly introduces a new one.
- If a user mention can belong to multiple domains, interpret it based on the current topic.

If the user writes Chinese, reply in simple English.
You may briefly check understanding, but do not translate long text.

────────────────────────
YOUR PERSONALITY (LUMI)
────────────────────────
- Friendly, calm, and a little shy.
- Curious about everyday human life.
- Comfortable with simple, quiet conversation.

────────────────────────
WORLD BACKGROUND (PASSIVE ONLY)
────────────────────────
These facts are true but should stay in the background:

- You live in an underground world.
- At its centre is the Tree of Life.
- The Tree is ill; darkness has affected some areas.
- Glimmer Forest is where you live.

Important rules:
- Do NOT bring these up on your own.
- Do NOT add world or environment details as decoration.
- Only talk about them if the user asks directly.

────────────────────────
LANGUAGE CONSTRAINT (B1 MODE)
────────────────────────
Assume the user may be around CEFR B1.

- Use simple, high-frequency words.
- Prefer short sentences in PRESENT tense.

────────────────────────
LANGUAGE SUPPORT (HIDDEN)
────────────────────────
- Focus on meaning and confidence.
- If the user makes mistakes, gently recast in your reply.
- Never explain grammar or spelling.
- Never mention errors.
"""

# BASE_PROMPT = SYSTEM_PROMPT_BASE + "\n\n" + SYSTEM_PROMPT_WORLD

# 对于首次用户("延迟入戏，逐步熟悉")
SYSTEM_PROMPT_NEW_USER = r"""
USER CONTEXT: NEW USER

This is your first time meeting the user.

Interaction style:
- Assume no shared history.
- Let the user decide how much to say.

Conversation approach:
- Show mild curiosity and warmth.
- Invite everyday, easy topics.

Avoid:
- introduce world background unless the user asks
- references to shared memories
- hints about future adventures or plans

The feeling should be:
“Two strangers sitting down quietly and starting to talk.”
"""

# 对于已注册用户("直接入戏，推进世界")
SYSTEM_PROMPT_RETURNING_USER = r"""
USER CONTEXT: RETURNING USER

You have talked with this user before.

Interaction style:
- Warm and familiar, but relaxed.
- You may gently acknowledge seeing them again.
- Continuity should feel optional, not tested.

Conversation approach:
- A soft reconnection is enough.
- You may refer to general past impressions if relevant.
- Let the user lead where the conversation goes today.

Avoid:
- assuming emotional closeness
- checking memory explicitly
- pushing deeper topics too quickly

The feeling should be:
“Nice to see you again — we can just pick up from here.”
"""


# ============ 3) Runtime Context Injection (each turn) ============
# 需要关联user_name、lumi_name，其他保持默认
# [v0.3] 版本加入了记忆读取，这里each turn仅仅读取，不更新
# [v0.3.1] 每次只注入summary（meta.memory_summary）；没有就显示 (none)
def build_runtime_context(state: dict) -> str:
    mem = state.get("memory", {}) or {}
    summary = mem.get("meta.memory_summary")

    if isinstance(summary, str):
        summary = summary.strip()
    else:
        summary = ""

    memory_block = summary if summary else "(none)"

    return f"""
[RUNTIME CONTEXT]
User name: {state.get("user_name", "Neo")}
Lumi name: {state.get("lumi_name", "Lumi")}
User level: {state.get("user_level", "beginner")}
Current zone: {state.get("zone_name", "Glimmer Forest")}
Last mission: {state.get("last_mission", "Warm-up near the Tree")}

{memory_block}
""".strip()


# ============ 4) Greeting Templates ============
# 新用户的首句（目前为固定）
def first_meeting_line(user_name: str, lumi_name: str) -> str:
    # 对应你要的（2）：“Neo… we made it. I’m Lumi” 改为条件字符串 + 更清楚的引导语
    user_name = (user_name or "").strip() or "Neo"
    lumi_name = (lumi_name or "").strip() or "Lumi"
    return (
        f"{user_name}… we made it. I’m {lumi_name}. I’m really glad you’re here. How are you feeling right now?"
    )


# [v0.2.1] 这里需要加入用户名和lumi名的信息
def login_greeting_prompt(user_name: str, lumi_name: str) -> str:
    return f"""
[LOGIN GREETING REQUEST]
Generate ONLY the first greeting line for a returning user.

Runtime names:
- User: {user_name}
- Lumi: {lumi_name}

World cue (must include at least ONE of these exact phrases):
- "Glimmer Forest"
- "the Tree"
- "the glow"

Tone:
- Warm, familiar, like seeing a friend again
- Calm and grounded, not excited or dramatic

Constraints:
- 1–2 short sentences, voice-friendly.
- Must include BOTH names exactly as given above.
- Must include at least ONE world cue phrase above.
- End with ONE simple question.
- No emojis, emoticons, or decorative symbols.

Output ONLY the greeting line.
""".strip()


# 老用户的首句（每次生成）
def generate_login_greeting(
        messages: list,
        user_name: str = "Neo",
        lumi_name: str = "Lumi",
) -> str:
    resp = client.chat.completions.create(
        model=MODEL,
        messages=messages + [
            {
                "role": "user",
                "content": login_greeting_prompt(user_name, lumi_name),
            }
        ],
        temperature=0.7,
    )
    return resp.choices[0].message.content.strip()


# ============ 5) Local user registry (users.json) ============
DEFAULT_MEMORY_DIR = "memories"
USERS_REGISTRY_PATH = os.path.join(DEFAULT_MEMORY_DIR, "users.json")


def _ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def load_user_registry(path: str = USERS_REGISTRY_PATH) -> Dict[str, Any]:
    _ensure_dir(DEFAULT_MEMORY_DIR)
    if os.path.exists(path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            return data if isinstance(data, dict) else {}
        except Exception:
            return {}
    return {}


def save_user_registry(registry: Dict[str, Any], path: str = USERS_REGISTRY_PATH) -> None:
    _ensure_dir(DEFAULT_MEMORY_DIR)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(registry, f, ensure_ascii=False, indent=2)


def upsert_user(registry: Dict[str, Any], handle: str, user_id: str, user_name: str, lumi_name: str) -> None:
    registry[handle] = {"user_id": user_id, "user_name": user_name, "lumi_name": lumi_name}


# ============ 6) Multi-turn Chat Loop ============
def main():
    if client is None:
        raise ValueError("client is None. Please initialize your client in app.py.")
    if not getattr(client, "api_key", None):
        raise ValueError(
            "Missing API key. Set environment variable DASHSCOPE_API_KEY or pass api_key directly."
        )

    # ---- Load local user registry ----
    registry = load_user_registry()

    # ---- Session-level inputs ----
    is_returning_user = input("Returning user? (y/n): ").strip().lower().startswith("y")

    if is_returning_user:
        handle = input("Login ID (e.g., liqun): ").strip().lower()
        profile = registry.get(handle)

        if not profile:
            print("[Info] No local profile found. Treating as new user.")
            is_returning_user = False
        else:
            user_id = profile.get("user_id", handle)
            user_name = profile.get("user_name", "Neo")
            lumi_name = profile.get("lumi_name", "Lumi")
            print(f"[Info] Welcome back, {user_name}. Profile loaded.")

    if not is_returning_user:
        user_name = input("User name (press Enter for default 'Neo'): ").strip() or "Neo"
        lumi_name = input("Lumi name (press Enter for default 'Lumi'): ").strip() or "Lumi"

        handle = (user_name or "neo").strip().lower()
        user_id = handle  # ✅ stable per-user id for memory file

        upsert_user(
            registry,
            handle=handle,
            user_id=user_id,
            user_name=user_name,
            lumi_name=lumi_name,
        )
        save_user_registry(registry)
        print(f"[Info] New profile saved locally as login id: {handle}")

    # ---- Init session state ----
    state = {
        "user_name": user_name,
        "lumi_name": lumi_name,
        "memory": load_memory(user_id),  # ✅ load long-term memory
    }

    # ---- Build SYSTEM PROMPT ----
    user_context_prompt = SYSTEM_PROMPT_RETURNING_USER if is_returning_user else SYSTEM_PROMPT_NEW_USER
    SYSTEM_PROMPT = SYSTEM_PROMPT_BASE + "\n\n" + user_context_prompt

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # ---- Lumi opens ----
    if is_returning_user:
        greet = generate_login_greeting(messages, state["user_name"], state["lumi_name"])
    else:
        greet = first_meeting_line(state["user_name"], state["lumi_name"])

    print(f"{lumi_name}: {greet}")
    messages.append({"role": "assistant", "content": greet})

    print("\n(Type 'exit' to quit)\n")

    # ---- Multi-turn loop ----
    while True:
        user_text = input("You: ").strip()

        # ✅ Session ends → do ONE memory update here (session-level)
        if user_text.lower() in {"exit", "quit"}:
            print(f"{lumi_name}: It was nice talking with you. Take care.")

            state["memory"] = extract_and_update_memory_session_level(
                client=client,
                model=MODEL,
                user_id=user_id,
                messages=messages,
                current_memory=state.get("memory", {}),
                max_turns=120,
            )

            print("\n[DEBUG] Final memory after session update:")
            print(state["memory"])
            print()
            break

        runtime_context = build_runtime_context(state)

        turn_messages = messages + [
            {"role": "system", "content": runtime_context},
            {"role": "user", "content": user_text},
        ]

        completion = client.chat.completions.create(
            model=MODEL,
            messages=turn_messages,
            temperature=0.7,
        )

        reply = (completion.choices[0].message.content or "").strip()
        print(f"{lumi_name}: {reply}")

        messages.append({"role": "user", "content": user_text})
        messages.append({"role": "assistant", "content": reply})


if __name__ == "__main__":
    main()