"""
Manages  in-memory conversation history for a session
Resets when app restarts

Structure:
    memory = {
    "summary":  str,      # compresses older convo turns
    "recent": list[dict]  # has {"role":"...", "content":"..."} of latest few turns
    }

When recent turns exceed certain number (RECENT_TURNS), the older turns are summarised and folded into 'summary', keeping the prompt size bounded without lossing track of earlier context.

What this file does:
1. initiate new empty memory 
2. summarize some older turns and updates the summary
3. add new turns of conversation in memory, and if the turn's number exceed the threshold, get summary
4. function to properly format the history or memory to pass to the prompt for LLM
"""

import os
from dotenv import load_dotenv
from google import genai
from config import LLM_MODEL
from logger import get_logger

load_dotenv()
logger = get_logger("memory")

llm = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

RECENT_TURNS = 6 # number of full (user+assistant) turn-pairs to keep
COMPRESS_BY = 3  # when overflow, fold this many oldest pairs into the summary

# 1. initiate new empty memory
def new_memory() -> dict:
    """Return a blank memory object"""
    return {"summary":"", "recent": []}

# 2. summarize some older turns and updates the summary
def summarise(existing_summary: str, turns_to_compress: list) -> str:
    """
    Ask the LLM to fold `turns_to_compress` into `existing_summary`.
    Returns the updated summary string.
    """
    try:
        turns_text = "\n".join(
            f"{'User' if t["role"] == 'user' else 'Assistant'}:{t['content']}"
            for t in turns_to_compress
        )

        prior_block = (
            f"Existing summary:\n{existing_summary}\n\n" if existing_summary else ""
        )

        prompt = (
            f"{prior_block}"
            f"New conversation turns to add to the summary: \n{turns_text}\n\n"
            "Write a concise but complete summary of the full conversation so far."
            "Preserve key facts, questions asked, and answers given."
            "Return ONLY the summary, nothing else."
        )

        response = llm.models.generate_content(
            model=LLM_MODEL,
            contents=prompt
        )

        new_summary = response.text.strip()
        logger.info("summarise: compressed %d turns into summary", len(turns_to_compress))
        return new_summary
    
    except Exception as e:
        logger.warning("summarise: LLM call failed, keeping existing summary: %s", e)
        return existing_summary  # don't lose old summary if compression fails


# 3. add new turns of conversation in memory, and if the turn's number exceed the threshold, get summary
def add_turn(memory: dict, role:str, content:str) -> dict:
    """
    Append one message to memory.
    If recent turns exceed the limit, compress the oldest ones into the summary.
    Always returns the updated memory dict.
    """

    memory["recent"].append({"role":role, "content":content})

    max_messages = RECENT_TURNS * 2 # pairs to individual messages
    if len(memory["recent"]) > max_messages:
        compress_count = COMPRESS_BY * 2 # messages to compress
        to_compress = memory["recent"][:compress_count] # take the starting compress_count messages 
        memory["recent"] = memory["recent"][compress_count:] # reassign the recent turns, starting from the compress_count nth turn to latest turn
        logger.info("add_turn: compressing %d messages into summary", compress_count)
        memory["summary"] = summarise(memory["summary"], to_compress)

    return memory

# 4. function to properly format the history or memory to pass to the prompt for LLM
def format_history_for_prompt(memory: dict) -> str:
    """
    Renders the full memory (summary + recent turns) as a plain-text block
    ready to inject into the LLM prompt.
    Returns an empty string if memory is empty.
    """  

    parts = []
    if memory.get("summary"):
        parts.append(f"Summary of earlier conversation:\n{memory['summary']}")

    if memory.get("recent"):
        recent_lines = []
        for turn in memory["recent"]:
            label = "User" if turn["role"] == "user" else "Assistant"
            recent_lines.append(f"{label}:{turn['content']}")
    parts.append("Recent conversation:\n" + "\n".join(recent_lines))

    return "\n\n".join(parts)