from datetime import datetime, timezone
from typing import Dict, Iterable, List, Optional

from memory_policy import evaluate_memory_write
from memory_store import MemoryStore
from schemas import MemoryEvent, RetrievedMemory


class MemoryAugmentedAgent:
    def __init__(self, memory_store: MemoryStore, audit_log: Optional[Iterable[dict]] = None):
        self.memory_store = memory_store
        self.audit_log = list(audit_log or [])

    def compress_context(self, retrieved_memories: List[RetrievedMemory]) -> str:
        if not retrieved_memories:
            return "No relevant project memories were retrieved."

        semantic = [item.memory.content for item in retrieved_memories if item.memory.memory_type == "semantic"]
        episodic = [item.memory.content for item in retrieved_memories if item.memory.memory_type == "episodic"]
        parts: List[str] = []

        if semantic:
            parts.append("Stable project context:")
            parts.extend(f"- {memory}" for memory in semantic)
        if episodic:
            if parts:
                parts.append("")
            parts.append("Relevant prior decisions and events:")
            parts.extend(f"- {memory}" for memory in episodic)

        return "\n".join(parts)

    def generate_response(self, query: str, retrieved_memories: List[RetrievedMemory]) -> str:
        if not retrieved_memories:
            return (
                "I do not have enough remembered project context to answer this from memory. "
                "The demo will not invent a project memory that was not retrieved."
            )

        lines = [
            "Here is the most relevant remembered project context for this request:",
            "",
        ]
        for item in retrieved_memories[:4]:
            lines.append(f"- {item.memory.content}")

        lines.extend(
            [
                "",
                "Memory boundary: this answer is derived only from the session memories shown in the inspection panel.",
            ]
        )
        return "\n".join(lines)

    def log_run(self, result: Dict) -> dict:
        retrieved_memory_ids = [item.memory.memory_id for item in result["retrieved_memories"]]
        decision = result["write_decision"]
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "query_length": len(result["query"]),
            "retrieved_memory_ids": retrieved_memory_ids,
            "write_outcome": decision.outcome,
            "write_reason": decision.reason,
            "policy_flags": decision.policy_flags,
            "memory_count_after_run": len(self.memory_store.memories),
        }
        self.audit_log.append(log_entry)
        return log_entry

    @staticmethod
    def _frame(event: MemoryEvent, result: Dict) -> Dict:
        return {"event": event, "result": dict(result)}

    def run_iter(self, query: str, top_k: int = 4, save_new_memory: bool = True):
        clean_query = (query or "").strip()
        result: Dict = {
            "query": clean_query,
            "retrieved_memories": [],
            "compressed_context": "",
            "agent_response": "",
            "write_decision": None,
            "saved_memory": None,
            "log_entry": None,
        }

        yield self._frame(
            MemoryEvent(stage="query_received", message="Request received. The agent is deciding which memories are relevant."),
            result,
        )

        retrieved = self.memory_store.retrieve(clean_query, top_k=top_k)
        result["retrieved_memories"] = retrieved
        yield self._frame(
            MemoryEvent(
                stage="memory_retrieved",
                message=f"Retrieved {len(retrieved)} ranked memories from this session's memory store.",
                memory_ids=[item.memory.memory_id for item in retrieved],
            ),
            result,
        )

        compressed_context = self.compress_context(retrieved)
        result["compressed_context"] = compressed_context
        yield self._frame(
            MemoryEvent(stage="context_compressed", message="Separated stable context from prior events and compressed the working context."),
            result,
        )

        write_decision = evaluate_memory_write(clean_query)
        result["write_decision"] = write_decision
        yield self._frame(
            MemoryEvent(
                stage="write_policy_evaluated",
                message=f"Memory write policy returned: {write_decision.outcome.upper()}.",
                decision=write_decision.outcome,
            ),
            result,
        )

        saved_memory = None
        if save_new_memory and write_decision.outcome == "save":
            saved_memory = self.memory_store.save_memory_from_decision(write_decision)
            result["saved_memory"] = saved_memory
            yield self._frame(
                MemoryEvent(
                    stage="memory_saved",
                    message="The approved project memory was added to this browser session only.",
                    memory_ids=[saved_memory.memory_id] if saved_memory else [],
                    decision="save",
                ),
                result,
            )
        elif write_decision.outcome == "block":
            yield self._frame(
                MemoryEvent(
                    stage="memory_blocked",
                    message="The application blocked this content from becoming long-term memory.",
                    decision="block",
                ),
                result,
            )
        else:
            yield self._frame(
                MemoryEvent(
                    stage="memory_skipped",
                    message=(
                        "Memory storage is disabled for this run."
                        if not save_new_memory and write_decision.outcome == "save"
                        else "The interaction was useful for this turn but was not saved as durable memory."
                    ),
                    decision="skip",
                ),
                result,
            )

        result["agent_response"] = self.generate_response(clean_query, retrieved)
        yield self._frame(
            MemoryEvent(stage="response_ready", message="Memory-grounded response is ready."),
            result,
        )

        result["log_entry"] = self.log_run(result)
        yield self._frame(
            MemoryEvent(stage="audit_recorded", message="Recorded control metadata without storing the raw query in the audit log."),
            result,
        )

    def run(self, query: str, top_k: int = 4, save_new_memory: bool = True):
        last_frame = None
        for frame in self.run_iter(query=query, top_k=top_k, save_new_memory=save_new_memory):
            last_frame = frame
        return last_frame["result"] if last_frame else None