import asyncio
import json
import logging
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional

from litellm import acompletion

from nexus_agent.models.memory import MemoryCategory, MemoryItem

logger = logging.getLogger(__name__)

EXTRACTION_PROMPT = """You are a memory extraction assistant. Analyze the following conversation exchange and extract key facts, preferences, patterns, or context worth remembering for future conversations.

## Existing Memories
{existing_memories}

## Rules
1. Only extract NEW information not already covered by existing memories.
2. Each memory should be a single, concise sentence.
3. Categorize each memory as one of: preference, context, pattern, fact
   - preference: user likes/dislikes, style preferences, language preferences
   - context: project context, work environment, ongoing tasks
   - pattern: recurring behaviors or request patterns
   - fact: personal facts, technical facts, domain knowledge
4. Return a JSON array of objects with "content" and "category" fields.
5. If there is nothing new worth remembering, return an empty array [].
6. Write memories in the same language the user is using.

## Conversation
User: {user_message}
Assistant: {assistant_message}

## Response (JSON array only, no markdown fencing):"""

COMPRESSION_PROMPT = """You are a memory compression assistant. Your job is to merge similar or overlapping memories to reduce the total count while preserving all important information.

## Current Memories (JSON)
{memories_json}

## Rules
1. Merge memories that are similar, overlapping, or about the same topic into a single combined memory.
2. Keep memories that are unique and unrelated to others as-is.
3. Each output item must include "source_ids" (array of original memory IDs that were merged into it), "content" (the merged text), and "category" (one of: preference, context, pattern, fact).
4. Every original memory ID must appear in exactly one output item's source_ids.
5. If a memory is not merged with anything, still include it with its single source_id.
6. Aim for at least 20% reduction in total count.
7. Write merged memories in the same language as the originals.

## Response (JSON array only, no markdown fencing):"""

_COMPRESSION_COOLDOWN = 300  # 5 minutes


class MemoryManager:
    def __init__(self):
        self._memories: Dict[str, MemoryItem] = {}
        self._config_path: Optional[Path] = None
        self._compression_lock = asyncio.Lock()
        self._last_compressed_at: Optional[float] = None

    def load_config(self, config_path: str) -> None:
        path = Path(config_path)
        if not path.is_absolute():
            from nexus_agent.config import get_config_path
            path = get_config_path(config_path)
        self._config_path = path

        if not path.exists():
            path.write_text(json.dumps({"memories": []}, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
            self._memories = {}
            return

        try:
            data = json.loads(path.read_text(encoding="utf-8"))
            for item in data.get("memories", []):
                mem = MemoryItem(**item)
                self._memories[mem.id] = mem
            logger.info(f"Loaded {len(self._memories)} memories from {path}")
        except Exception as e:
            logger.warning(f"Failed to parse memories: {e}, starting fresh")
            self._memories = {}

    def _save_config(self) -> None:
        if not self._config_path:
            return
        data = {"memories": [m.model_dump() for m in self._memories.values()]}
        self._config_path.write_text(
            json.dumps(data, indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )

    # --- CRUD ---

    def get_all(self) -> List[MemoryItem]:
        return sorted(self._memories.values(), key=lambda m: m.created_at, reverse=True)

    def get(self, memory_id: str) -> Optional[MemoryItem]:
        return self._memories.get(memory_id)

    def create(self, content: str, category: MemoryCategory = "fact") -> MemoryItem:
        now = datetime.now(timezone.utc).isoformat()
        mem = MemoryItem(
            id=uuid.uuid4().hex[:12],
            content=content,
            category=category,
            created_at=now,
            updated_at=now,
        )
        self._memories[mem.id] = mem
        self._save_config()
        logger.info(f"Created memory [{category}]: {content[:50]}")
        return mem

    def update(self, memory_id: str, **kwargs) -> Optional[MemoryItem]:
        mem = self._memories.get(memory_id)
        if not mem:
            return None
        data = mem.model_dump()
        for k, v in kwargs.items():
            if v is not None:
                data[k] = v
        data["updated_at"] = datetime.now(timezone.utc).isoformat()
        self._memories[memory_id] = MemoryItem(**data)
        self._save_config()
        return self._memories[memory_id]

    def delete(self, memory_id: str) -> bool:
        if memory_id not in self._memories:
            return False
        self._memories.pop(memory_id)
        self._save_config()
        logger.info(f"Deleted memory: {memory_id}")
        return True

    def clear_all(self) -> int:
        count = len(self._memories)
        self._memories.clear()
        self._save_config()
        logger.info(f"Cleared all {count} memories")
        return count

    # --- System Prompt ---

    def build_memory_prompt(self) -> str:
        """Build a memory section for the system prompt."""
        from nexus_agent.core.settings_manager import settings_manager

        memory_settings = settings_manager.settings.memory
        if not memory_settings.enabled or not self._memories:
            return ""

        lines = []
        for mem in self.get_all():
            lines.append(f"- [{mem.category}] {mem.content}")
            # Rough token estimate: ~4 chars per token
            if len("\n".join(lines)) > memory_settings.max_injection_tokens * 4:
                break

        if not lines:
            return ""

        return (
            "## Long-term Memory\n"
            "아래는 이전 대화에서 기억된 사실과 선호도입니다:\n\n"
            + "\n".join(lines)
        )

    # --- Compression & Eviction ---

    def _evict_oldest(self) -> Optional[str]:
        """Delete the oldest non-pinned memory by created_at. Returns the deleted ID."""
        candidates = [m for m in self._memories.values() if not m.is_pinned]
        if not candidates:
            return None
        oldest = min(candidates, key=lambda m: m.created_at)
        self.delete(oldest.id)
        logger.info(f"Auto-replaced oldest memory: {oldest.id} ({oldest.content[:40]})")
        return oldest.id

    async def _compress_memories(self) -> int:
        """Merge similar memories via LLM. Returns number of freed slots."""
        async with self._compression_lock:
            # Cooldown check
            now = time.monotonic()
            if self._last_compressed_at and (now - self._last_compressed_at) < _COMPRESSION_COOLDOWN:
                logger.debug("Memory compression skipped: cooldown active")
                return 0

            # Only compress non-pinned memories
            unpinned = [m for m in self._memories.values() if not m.is_pinned]

            # Need at least 4 non-pinned memories to compress
            if len(unpinned) < 4:
                return 0

            from nexus_agent.core.settings_manager import settings_manager
            llm = settings_manager.llm

            all_mems = sorted(unpinned, key=lambda m: m.created_at, reverse=True)
            memories_data = [
                {"id": m.id, "content": m.content, "category": m.category}
                for m in all_mems
            ]
            original_count = len(memories_data)
            original_ids = {m["id"] for m in memories_data}

            prompt = COMPRESSION_PROMPT.format(
                memories_json=json.dumps(memories_data, ensure_ascii=False)
            )

            try:
                from nexus_agent.core.llm import LLMClient
                api_key = llm.api_key or LLMClient._resolve_api_key(llm.model)

                kwargs = {
                    "model": llm.model,
                    "api_key": api_key,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.2,
                    "max_tokens": 4096,
                }
                if llm.api_base:
                    kwargs["api_base"] = llm.api_base

                response = await acompletion(**kwargs)
                raw = (response.choices[0].message.content or "").strip()

                if not raw:
                    return 0

                # Strip markdown code fencing if present
                if raw.startswith("```"):
                    raw = raw.split("\n", 1)[1] if "\n" in raw else raw[3:]
                    if raw.endswith("```"):
                        raw = raw[:-3].strip()

                compressed = json.loads(raw)
                if not isinstance(compressed, list):
                    logger.warning("Memory compression: LLM returned non-list")
                    return 0

                # Validate: every original ID must appear exactly once
                seen_ids: set[str] = set()
                for item in compressed:
                    source_ids = item.get("source_ids", [])
                    for sid in source_ids:
                        if sid in seen_ids:
                            logger.warning(f"Memory compression aborted: duplicate source_id {sid}")
                            return 0
                        seen_ids.add(sid)

                if seen_ids != original_ids:
                    missing = original_ids - seen_ids
                    extra = seen_ids - original_ids
                    logger.warning(
                        f"Memory compression aborted: ID mismatch "
                        f"(missing={missing}, extra={extra})"
                    )
                    return 0

                # Apply merges
                freed = 0
                for item in compressed:
                    source_ids = item.get("source_ids", [])
                    if len(source_ids) <= 1:
                        # No merge — keep original as-is
                        continue

                    content = item.get("content", "").strip()
                    category = item.get("category", "fact")
                    if not content:
                        continue
                    if category not in ("preference", "context", "pattern", "fact"):
                        category = "fact"

                    # Delete originals
                    for sid in source_ids:
                        if sid in self._memories:
                            del self._memories[sid]

                    # Create merged memory
                    now_ts = datetime.now(timezone.utc).isoformat()
                    mem = MemoryItem(
                        id=uuid.uuid4().hex[:12],
                        content=content,
                        category=category,
                        created_at=now_ts,
                        updated_at=now_ts,
                    )
                    self._memories[mem.id] = mem
                    freed += len(source_ids) - 1

                if freed > 0:
                    self._save_config()
                    new_count = len(self._memories)
                    logger.info(
                        f"Memory compression complete: {original_count} → {new_count} "
                        f"(freed {freed} slots)"
                    )

                self._last_compressed_at = time.monotonic()
                return freed

            except Exception as e:
                logger.warning(f"Memory compression failed: {e}")
                return 0

    # --- Extraction ---

    async def extract_and_save(self, user_message: str, assistant_message: str) -> List[MemoryItem]:
        """Extract memories from a conversation turn using LLM."""
        from nexus_agent.core.settings_manager import settings_manager

        memory_settings = settings_manager.settings.memory
        if not memory_settings.enabled:
            return []

        # Phase 1: Compress if capacity threshold exceeded
        capacity_ratio = len(self._memories) / max(memory_settings.max_memories, 1)
        if capacity_ratio >= memory_settings.compression_threshold:
            freed = await self._compress_memories()
            if freed > 0:
                logger.info(f"Compression freed {freed} slots before extraction")

        # Build existing memories context
        existing = "\n".join(
            f"- [{m.category}] {m.content}" for m in self.get_all()
        ) or "(none)"

        prompt = EXTRACTION_PROMPT.format(
            existing_memories=existing,
            user_message=user_message,
            assistant_message=assistant_message,
        )

        try:
            llm = settings_manager.llm
            from nexus_agent.core.llm import LLMClient
            api_key = llm.api_key or LLMClient._resolve_api_key(llm.model)

            kwargs = {
                "model": llm.model,
                "api_key": api_key,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.3,
                "max_tokens": 1024,
            }
            if llm.api_base:
                kwargs["api_base"] = llm.api_base

            response = await acompletion(**kwargs)
            raw = (response.choices[0].message.content or "").strip()

            if not raw:
                return []

            # Strip markdown code fencing if present
            if raw.startswith("```"):
                raw = raw.split("\n", 1)[1] if "\n" in raw else raw[3:]
                if raw.endswith("```"):
                    raw = raw[:-3].strip()

            extracted = json.loads(raw)
            if not isinstance(extracted, list):
                return []

            created = []
            for item in extracted:
                content = item.get("content", "").strip()
                category = item.get("category", "fact")
                if not content:
                    continue
                if category not in ("preference", "context", "pattern", "fact"):
                    category = "fact"
                # Phase 3: Evict oldest if at capacity
                if len(self._memories) >= memory_settings.max_memories:
                    self._evict_oldest()
                mem = self.create(content, category)
                created.append(mem)

            if created:
                logger.info(f"Extracted {len(created)} new memories from conversation")
            return created

        except Exception as e:
            logger.warning(f"Memory extraction failed: {e}")
            return []


memory_manager = MemoryManager()
