"""Memory system for persistent agent memory."""

from __future__ import annotations

import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import TYPE_CHECKING

from loguru import logger

from asedclaw.utils.helpers import ensure_dir

if TYPE_CHECKING:
    from asedclaw.providers.base import LLMProvider
    from asedclaw.session.manager import Session


class MemoryStore:
    """Two-layer memory: MEMORY.md + append-only daily logs (memory/YYYY-MM-DD.md)."""

    def __init__(self, workspace: Path):
        self.memory_dir = ensure_dir(workspace / "memory")
        self.memory_file = self.memory_dir / "MEMORY.md"
        # Legacy file kept only for migration reads.
        self.legacy_history_file = self.memory_dir / "HISTORY.md"

    def read_long_term(self) -> str:
        if self.memory_file.exists():
            return self.memory_file.read_text(encoding="utf-8")
        return ""

    def write_long_term(self, content: str) -> None:
        self.memory_file.write_text(content, encoding="utf-8")

    def get_daily_file(self, when: datetime | None = None) -> Path:
        ts = when or datetime.now()
        return self.memory_dir / f"{ts:%Y-%m-%d}.md"

    def append_daily(self, entry: str, when: datetime | None = None) -> None:
        daily_file = self.get_daily_file(when)
        with open(daily_file, "a", encoding="utf-8") as f:
            f.write(entry.rstrip() + "\n\n")

    def read_recent_daily(self, days: int = 2, max_chars: int = 6000) -> str:
        """Read recent daily logs (today backwards). Includes legacy HISTORY.md as fallback."""
        chunks: list[str] = []
        now = datetime.now()
        for offset in range(max(1, days)):
            fp = self.get_daily_file(now - timedelta(days=offset))
            if fp.exists():
                content = fp.read_text(encoding="utf-8").strip()
                if content:
                    chunks.append(f"### {fp.name}\n{content}")

        if not chunks and self.legacy_history_file.exists():
            legacy = self.legacy_history_file.read_text(encoding="utf-8").strip()
            if legacy:
                chunks.append(f"### HISTORY.md (legacy)\n{legacy}")

        merged = "\n\n".join(chunks)
        if len(merged) > max_chars:
            return merged[:max_chars] + "\n\n...(truncated)"
        return merged

    def get_memory_context(self) -> str:
        parts: list[str] = []
        long_term = self.read_long_term()
        if long_term:
            parts.append(f"## Long-term Memory\n{long_term}")

        recent = self.read_recent_daily(days=2)
        if recent:
            parts.append(f"## Recent Daily Memory\n{recent}")

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

    async def consolidate(
        self,
        session: Session,
        provider: LLMProvider,
        model: str,
        *,
        archive_all: bool = False,
        memory_window: int = 50,
    ) -> bool:
        """Consolidate old messages into MEMORY.md + daily logs via LLM tool call.

        Returns True on success (including no-op), False on failure.
        """
        if archive_all:
            old_messages = session.messages
            keep_count = 0
            logger.info("Memory consolidation (archive_all): {} messages", len(session.messages))
        else:
            keep_count = memory_window // 2
            if len(session.messages) <= keep_count:
                return True
            if len(session.messages) - session.last_consolidated <= 0:
                return True
            old_messages = session.messages[session.last_consolidated:-keep_count]
            if not old_messages:
                return True
            logger.info("Memory consolidation: {} to consolidate, {} keep", len(old_messages), keep_count)

        lines = []
        for m in old_messages:
            if not m.get("content"):
                continue
            tools = f" [tools: {', '.join(m['tools_used'])}]" if m.get("tools_used") else ""
            lines.append(f"[{m.get('timestamp', '?')[:16]}] {m['role'].upper()}{tools}: {m['content']}")

        current_memory = self.read_long_term()
        prompt = f"""Process this conversation and call the save_memory tool with your consolidation.
The daily_entry should be append-only notes for today's file (memory/YYYY-MM-DD.md).
The memory_update should be the full updated content for MEMORY.md.

## Safety Rules
1. daily_entry is append-only — if the daily file already exists, your entry will be appended.
2. MEMORY.md content must go through the memory_update return value only — do NOT write it via other tools.
3. Do NOT modify SOUL.md, TOOLS.md, or AGENTS.md.

## Current Long-term Memory
{current_memory or "(empty)"}

## Conversation to Process
{chr(10).join(lines)}"""

        try:
            args = await self._consolidate_with_agno(
                prompt=prompt,
                model=model,
                api_key=getattr(provider, "api_key", None),
                api_base=getattr(provider, "api_base", None),
            )
            if args is None:
                logger.warning("Memory consolidation: agno did not produce save_memory payload")
                return False

            if not isinstance(args, dict):
                logger.warning("Memory consolidation: unexpected arguments type {}", type(args).__name__)
                return False

            if entry := args.get("daily_entry"):
                if not isinstance(entry, str):
                    entry = json.dumps(entry, ensure_ascii=False)
                self.append_daily(entry)
            elif entry := args.get("history_entry"):
                # Compatibility with old payload name.
                if not isinstance(entry, str):
                    entry = json.dumps(entry, ensure_ascii=False)
                self.append_daily(entry)
            if update := args.get("memory_update"):
                if not isinstance(update, str):
                    update = json.dumps(update, ensure_ascii=False)
                if update != current_memory:
                    self.write_long_term(update)

            session.last_consolidated = 0 if archive_all else len(session.messages) - keep_count
            logger.info("Memory consolidation done: {} messages, last_consolidated={}", len(session.messages), session.last_consolidated)
            return True
        except Exception:
            logger.exception("Memory consolidation failed")
            return False

    async def _consolidate_with_agno(
        self,
        prompt: str,
        model: str,
        api_key: str | None,
        api_base: str | None,
    ) -> dict | None:
        """Run memory consolidation via agno agent tool call."""
        try:
            from agno.agent.agent import Agent
            from agno.models.openai import OpenAILike
            from agno.tools import tool
        except Exception:
            logger.exception("Agno memory runtime unavailable")
            return None

        payload: dict[str, str] = {}

        @tool(name="save_memory", description="Save consolidated memory output.")
        async def save_memory(daily_entry: str, memory_update: str) -> str:
            payload["daily_entry"] = daily_entry
            payload["memory_update"] = memory_update
            return "ok"

        model_kwargs: dict[str, str] = {"id": model}
        if api_key:
            model_kwargs["api_key"] = api_key
        if api_base:
            model_kwargs["base_url"] = api_base

        agent = Agent(
            model=OpenAILike(**model_kwargs),
            tools=[save_memory],
            instructions=(
                "You are a memory consolidation agent. "
                "Call save_memory(daily_entry, memory_update) exactly once. "
                "daily_entry must be append-only notes for today's daily log file. "
                "Do NOT modify SOUL.md, TOOLS.md, or AGENTS.md. "
                "MEMORY.md content goes through memory_update only."
            ),
            stream=False,
        )
        await agent.arun(prompt, stream=False)
        return payload or None
