main.py

"""
main.py — Thin CLI orchestrator for the TOWN “tablet‑core” demo
----------------------------------------------------------------
• Usage:  python main.py <npc_name>
• Type “quit” to exit.
"""

from __future__ import annotations

import logging
import sys
import json
import re
from pathlib import Path

from town_core.state_engine import GameState
from town_core.memory_engine import MemoryEngine
from town_core.prompt_engine import PromptEngine
from town_core.llm_engine import LLMEngine
from town_core.image_engine import ImageEngine
from town_core.embed_engine import EmbedEngine

# ----------------------------------------------------------------------------
# Template name
# ----------------------------------------------------------------------------
NPC_CHAT_TEMPLATE = "npc_chat"

# ----------------------------------------------------------------------------
# RelationshipSentences Loader
# ----------------------------------------------------------------------------
class RelationshipSentences:
    def __init__(self, filepath):
        with open(filepath, "r", encoding="utf-8") as f:
            self.data = json.load(f)

    def get_sentence(self, stat_name: str, value: int) -> str:
        clamped_value = max(-10, min(10, int(round(value))))
        return self.data.get(stat_name, {}).get(str(clamped_value), "No information available.")

# ----------------------------------------------------------------------------
# Helper: create placeholder sheet on first run
# ----------------------------------------------------------------------------
def bootstrap_npc(npc_name: str, mem: MemoryEngine) -> None:
    sheet_path = Path(f"npc_data/{npc_name}/character_sheet.txt")
    if sheet_path.exists():
        return
    text = (
        f"{npc_name.replace('_', ' ').title()} is a placeholder NPC for the "
        "lean‑core demo.  Provide a richer sheet later."
    )
    mem.ensure_sheet(npc_name, text)
    logging.info(f"[Setup] Created default sheet for {npc_name}")

# ----------------------------------------------------------------------------
def extract_response_text(full_reply: str) -> tuple[str, str]:
    # Extract think content
    think_start = full_reply.find("<think>") + len("<think>")
    think_end = full_reply.find("</think>")
    think_text = full_reply[think_start:think_end].strip() if think_start != -1 and think_end != -1 else ""
    
    # Extract response content
    response_start = full_reply.find("<response>") + len("<response>")
    response_end = full_reply.find("</response>")
    if response_start != -1 and response_end != -1:
        response_text = full_reply[response_start:response_end].strip()
    else:
        # Fallback: Remove think tags and use remaining text
        response_text = full_reply
        if think_text:
            response_text = response_text.replace(f"<think>{think_text}</think>", "").strip()
    
    return response_text, think_text

# ----------------------------------------------------------------------------
def main():
    if len(sys.argv) < 2:
        logging.error("Usage: python main.py <npc_name>")
        sys.exit(1)
    npc_name = sys.argv[1].lower()

    # ---------- singletons ----------
    gs = GameState()
    mem = MemoryEngine()
    pe = PromptEngine()
    llm = LLMEngine()
    cfg = llm.cfg
    img_engine = ImageEngine(cfg, npc_name) if cfg.get("image_model") else None
    rel_sentences = RelationshipSentences("town_core/relationship_sentences.json")
    embed_engine = EmbedEngine()

    # ---------- static rules ----------
    FULL_SYS = Path("FULL_SYSTEM_PROMPT.txt").read_text(encoding="utf-8")
    FULL_ASST = Path("FULL_ASSISTANT_PROMPT.txt").read_text(encoding="utf-8")

    # ---------- initial state ----------
    gs.npc_name = npc_name
    gs.location = "default"
    bootstrap_npc(npc_name, mem)

    # path for rolling history file
    hist_path = Path(f"npc_data/{npc_name}/conversation_history.txt")
    hist_path.parent.mkdir(parents=True, exist_ok=True)
    knowledge_path = Path(f"npc_data/{npc_name}/summoned_knowledge.txt")

    # Player description from config, with fallback
    player_description = cfg.get("player_description", "A curious adventurer with short brown hair, wearing a leather jacket and carrying a worn satchel.")

    logging.info(f"--- TALKING TO {npc_name.upper()} (type 'quit' to exit) ---")

    just_summarized = False

    def extract_knowledge_tags(response: str) -> List[str]:
        tags = []
        start = 0
        while True:
            start_tag = response.find("<summon_knowledge>", start)
            if start_tag == -1:
                break
            end_tag = response.find("</summon_knowledge>", start_tag)
            if end_tag == -1:
                break
            tag_content = response[start_tag + 17:end_tag].strip()
            tags.extend([t.strip() for t in tag_content.split(",")[:5]])
            start = end_tag + 18
        return tags

    while True:
        player_line = input("YOU > ").strip()
        if player_line.lower() in {"quit", "exit"}:
            logging.info("Exiting…")
            break
        if not player_line:
            continue

        gs.player_line = player_line
        gs.step_stage()

        ctx = mem.fetch_context(npc_name, player_line, k_per_bucket=4)

        ctx["trust_sentence"] = rel_sentences.get_sentence("trust", ctx["metrics"]["trust"])
        ctx["affection_sentence"] = rel_sentences.get_sentence("affection", ctx["metrics"]["affection"])
        ctx["respect_sentence"] = rel_sentences.get_sentence("respect", ctx["metrics"]["respect"])

        # Load summoned knowledge if exists
        ctx["summoned_knowledge"] = knowledge_path.read_text(encoding="utf-8") if knowledge_path.exists() else ""
        # Blank knowledge file after injection
        if knowledge_path.exists():
            knowledge_path.write_text("", encoding="utf-8")

        prompt = pe.render(
            NPC_CHAT_TEMPLATE,
            npc_name=npc_name,
            system_rules=FULL_SYS,
            assistant_rules=FULL_ASST,
            character_sheet=ctx["character_sheet"][0] if ctx["character_sheet"] else "No character sheet available.",
            memory_chunks=ctx["memory_chunks"],
            character_memories=ctx["character_memories"],
            location_memories=ctx["location_memories"],
            lore_memories=ctx["lore_memories"],
            summary_memories=ctx["summary_memories"],
            summoned_knowledge=ctx["summoned_knowledge"],
            player_line=player_line,
            trust_sentence=ctx["trust_sentence"],
            affection_sentence=ctx["affection_sentence"],
            respect_sentence=ctx["respect_sentence"]
        )

        full_reply = llm.chat(prompt)
        gs.mark_llm_call()

        response_text, think_text = extract_response_text(full_reply)

        # Extract and process summon_knowledge tags
        knowledge_tags = extract_knowledge_tags(full_reply)
        if knowledge_tags:
            knowledge_results = []
            for tag in knowledge_tags[:5]:  # Limit to 5 terms
                hits = embed_engine.search(tag, "knowledge_system", k=3)
                if hits:
                    knowledge_results.append(f"{tag}:")
                    knowledge_results.extend([hit["text"] for hit in hits if "text" in hit])
            if knowledge_results:
                knowledge_path.write_text("\n".join(knowledge_results), encoding="utf-8")
                logger.info("Saved knowledge to %s", knowledge_path)

        mem.update_metrics_from_response(npc_name, full_reply)

        print(f"{npc_name.upper()} > {response_text}\n")

        with hist_path.open("a", encoding="utf-8") as fh:
            fh.write(f"PLAYER: {player_line}\n{npc_name.upper()}: {response_text}\n")

        idx = len(list(Path("llm_logs").glob("llmIO_*.txt"))) + 1
        with open(f"llm_logs/llmIO_{idx:03d}.txt", "w", encoding="utf-8") as fh:
            fh.write(f"=== PROMPT ===\n{prompt}\n\n=== THINK ===\n{think_text}\n\n=== RESPONSE ===\n{full_reply}")

        if hist_path.exists():
            history = hist_path.read_text(encoding="utf-8")
            if len(history) > mem.HISTORY_WINDOW:
                mem.summarize_history(npc_name, history)
                just_summarized = True
            else:
                just_summarized = False

            if just_summarized and img_engine:
                summary_path = Path(f"npc_data/{npc_name}/conversation_summary.txt")
                summary = summary_path.read_text(encoding="utf-8").split("\n")[-1] if summary_path.exists() else ""
                character_description = ctx["character_sheet"][0] if ctx["character_sheet"] else ""
                img_engine.generate_and_display(
                    summary=summary,
                    location=gs.location,
                    character_description=character_description,
                    player_description=player_description
                )

if __name__ == "__main__":
    main()