"""Conversation memory with AI-powered compression and persistence.

When the conversation gets too long for the model's context window,
old messages are automatically summarized into a compact digest.
The summary + recent messages always fit within the token budget.
Conversations are saved to disk so they persist across app restarts.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
import re
import threading
import time
import uuid
from collections.abc import Callable

from .utils.atomic_io import atomic_write_json, safe_load_json

log = logging.getLogger(__name__)

# Sentinel for detecting corrupt JSON files
_CORRUPT = object()

# Token estimation weights per character class.
# CJK ideographs tokenize at ~1.5 tokens per character,
# Cyrillic at ~1.1 tokens per character, Latin/ASCII at ~0.25 per character.
_WEIGHT_CJK = 1.5
_WEIGHT_CYRILLIC = 1.1
_WEIGHT_DEFAULT = 0.25


def _is_cjk(ch: str) -> bool:
    """Return True if the character is a CJK ideograph."""
    cp = ord(ch)
    return (
        0x4E00 <= cp <= 0x9FFF  # CJK Unified Ideographs
        or 0x3400 <= cp <= 0x4DBF  # CJK Unified Ext A
        or 0xF900 <= cp <= 0xFAFF  # CJK Compatibility Ideographs
        or 0x20000 <= cp <= 0x2A6DF  # CJK Ext B
    )


def estimate_tokens(text: str) -> int:
    """Estimate token count with CJK/Cyrillic-aware weighting."""
    total = 0.0
    for ch in text:
        if _is_cjk(ch):
            total += _WEIGHT_CJK
        elif 0x0400 <= ord(ch) <= 0x04FF:
            total += _WEIGHT_CYRILLIC
        else:
            total += _WEIGHT_DEFAULT
    return max(1, int(total))


class ConversationMemory:
    """Maintains unlimited multi-turn conversation context.

    When the token budget is exceeded, older messages are compressed
    into a summary using an AI summarizer function.
    """

    def __init__(
        self,
        max_context_tokens: int = 80_000,
        keep_recent: int = 6,
        data_dir: str | None = None,
        max_messages: int = 100,
        max_tokens: int = 128_000,
    ) -> None:
        self.max_context_tokens = max_context_tokens
        self.keep_recent = keep_recent
        self._max_messages = max_messages
        self._max_tokens = max_tokens

        self._lock: threading.RLock = threading.RLock()
        self.messages: list[dict] = []
        self.system_prompt: str = ""
        self.summary: str = ""

        self._data_dir = data_dir or os.path.expanduser("~/.muravey-china/conversations")
        self._conversation_id: str | None = None
        self._created_at: str | None = None
        self._title: str | None = None

        self._total_messages_ever: int = 0
        self._compressions: int = 0

    # =========================================================================
    # Public API — messages
    # =========================================================================

    def set_system_prompt(self, prompt: str) -> None:
        self.system_prompt = prompt

    def add_user_message(self, text: str) -> None:
        """Add user message with automatic pruning."""
        with self._lock:
            self.messages.append({"role": "user", "content": text})
            self._total_messages_ever += 1
            self._prune_if_needed()

    def add_assistant_message(self, text: str) -> None:
        """Add assistant message with automatic pruning."""
        with self._lock:
            self.messages.append({"role": "assistant", "content": text})
            self._total_messages_ever += 1
            self._prune_if_needed()

    def get_messages(self) -> list[dict]:
        """Return conversation messages (without system prompt or summary)."""
        with self._lock:
            return list(self.messages)

    def get_messages_with_system(self) -> list[dict]:
        """Return full message list for the AI provider.

        Structure:
        1. System prompt (original)
        2. Summary of older conversation (if any)
        3. Recent messages
        """
        with self._lock:
            msgs = []
            if self.system_prompt:
                system_content = self.system_prompt
                if self.summary:
                    system_content += (
                        "\n\n--- PREVIOUS CONVERSATION CONTEXT ---\n"
                        f"{self.summary}\n"
                        "--- END CONTEXT ---\n"
                        "Continue the conversation naturally, using the context above."
                    )
                msgs.append({"role": "system", "content": system_content})
            elif self.summary:
                msgs.append(
                    {
                        "role": "system",
                        "content": (
                            f"Previous conversation context:\n{self.summary}\n\n"
                            "Continue the conversation naturally."
                        ),
                    },
                )
            msgs.extend(self.messages)
            return msgs

    def clear(self) -> None:
        """Start a new conversation."""
        with self._lock:
            if self.messages and self._conversation_id:
                self.save()
            self.messages = []
            self.system_prompt = ""
            self.summary = ""
            self._conversation_id = None
            self._created_at = None
            self._title = None
            self._total_messages_ever = 0
            self._compressions = 0
        log.debug("Conversation memory cleared")

    # =========================================================================
    # Properties
    # =========================================================================

    @property
    def has_history(self) -> bool:
        with self._lock:
            return len(self.messages) > 0

    @property
    def turn_count(self) -> int:
        with self._lock:
            return len([m for m in self.messages if m["role"] == "user"])

    @property
    def estimated_tokens(self) -> int:
        with self._lock:
            total = estimate_tokens(self.system_prompt)
            if self.summary:
                total += estimate_tokens(self.summary)
            for m in self.messages:
                total += estimate_tokens(m["content"])
            return total

    @property
    def needs_compression(self) -> bool:
        return self.estimated_tokens > self.max_context_tokens

    @property
    def conversation_id(self) -> str | None:
        return self._conversation_id

    @property
    def title(self) -> str | None:
        return self._title

    @title.setter
    def title(self, value: str) -> None:
        self._title = value

    # =========================================================================
    # Context Limits — automatic pruning
    # =========================================================================

    def _prune_if_needed(self) -> None:
        """Prune old messages if limits exceeded."""
        if len(self.messages) > self._max_messages:
            system_msgs = [m for m in self.messages if m["role"] == "system"]
            other_msgs = [m for m in self.messages if m["role"] != "system"]
            keep = max(1, self._max_messages - len(system_msgs))
            self.messages = system_msgs + other_msgs[-keep:]
            log.info("Pruned messages to %d", len(self.messages))

        total_tokens = sum(len(str(m.get("content", "")).split()) * 1.3 for m in self.messages)
        if total_tokens > self._max_tokens:
            while total_tokens > self._max_tokens and len(self.messages) > 2:
                for i, msg in enumerate(self.messages):
                    if msg["role"] != "system":
                        removed = self.messages.pop(i)
                        total_tokens -= len(str(removed.get("content", "")).split()) * 1.3
                        break
            log.info("Pruned by tokens: %.0f tokens", total_tokens)

    # =========================================================================
    # Compression — AI-powered summarization
    # =========================================================================

    def compress_if_needed(self, summarize_fn: Callable[[str], str]) -> None:
        """Compress old messages into a summary if token budget is exceeded.

        Args:
            summarize_fn: function(text) -> summary_text
        """
        if not self.needs_compression:
            return

        with self._lock:
            if len(self.messages) <= self.keep_recent:
                return

            split_idx = len(self.messages) - self.keep_recent
            old_messages = self.messages[:split_idx]
            self.messages = self.messages[split_idx:]

            parts = []
            if self.summary:
                parts.append(f"Previous summary:\n{self.summary}")
            for m in old_messages:
                role = "User" if m["role"] == "user" else "Assistant"
                parts.append(f"{role}: {m['content']}")
            text_to_summarize = "\n\n".join(parts)

        try:
            new_summary = summarize_fn(text_to_summarize)
            with self._lock:
                self.summary = new_summary.strip()
                self._compressions += 1
                log.info(
                    "Compressed: %d old messages -> summary (%d tokens), "
                    "kept %d recent, compressions: %d",
                    len(old_messages),
                    estimate_tokens(self.summary),
                    len(self.messages),
                    self._compressions,
                )
        except Exception as e:
            log.warning("Compression failed: %s. Dropping oldest messages.", e)
            with self._lock:
                self.messages = old_messages[-2:] + self.messages

    # =========================================================================
    # Persistence
    # =========================================================================

    @staticmethod
    def _is_safe_id(conversation_id: str) -> bool:
        return bool(conversation_id and re.fullmatch(r"[a-zA-Z0-9_-]+", conversation_id))

    def ensure_id(self) -> None:
        if not self._conversation_id:
            self._conversation_id = uuid.uuid4().hex[:12]
            self._created_at = time.strftime("%Y-%m-%d %H:%M:%S")

    def save(self) -> None:
        """Save current conversation to disk."""
        with self._lock:
            if not self.messages:
                return
            self.ensure_id()
            os.makedirs(self._data_dir, exist_ok=True)

            data = {
                "id": self._conversation_id,
                "title": self._title,
                "created_at": self._created_at,
                "updated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
                "system_prompt": self.system_prompt,
                "summary": self.summary,
                "messages": self.messages,
                "total_messages_ever": self._total_messages_ever,
                "compressions": self._compressions,
            }

            path = os.path.join(self._data_dir, f"{self._conversation_id}.json")
            try:
                atomic_write_json(path, data)
                log.debug(
                    "Saved conversation %s (%d messages)", self._conversation_id, len(self.messages)
                )
            except OSError as e:
                log.warning("Failed to save conversation: %s", e)

        self._save_last_id()

    def load(self, conversation_id: str) -> bool:
        """Load a conversation from disk. Returns True on success."""
        if not self._is_safe_id(conversation_id):
            log.warning("Invalid conversation ID rejected: %s", repr(conversation_id))
            return False
        path = os.path.join(self._data_dir, f"{conversation_id}.json")
        if not os.path.exists(path):
            log.warning("Conversation not found: %s", conversation_id)
            return False

        data = safe_load_json(path, default=_CORRUPT)
        if data is _CORRUPT:
            log.warning("Failed to load conversation: %s", conversation_id)
            return False

        with self._lock:
            self._conversation_id = data.get("id", conversation_id)
            self._created_at = data.get("created_at")
            self._title = data.get("title")
            self.system_prompt = data.get("system_prompt", "")
            self.summary = data.get("summary", "")
            self.messages = data.get("messages", [])
            self._total_messages_ever = data.get("total_messages_ever", len(self.messages))
            self._compressions = data.get("compressions", 0)

        log.info("Loaded conversation %s: %d messages", self._conversation_id, len(self.messages))
        return True

    def load_last(self) -> bool:
        """Load the most recent conversation. Returns True on success."""
        last_id = self._get_last_id()
        if last_id:
            return self.load(last_id)
        return False

    def list_conversations(self, limit: int = 20) -> list[dict]:
        """List recent conversations (metadata only)."""
        os.makedirs(self._data_dir, exist_ok=True)
        convos = []
        for filename in sorted(os.listdir(self._data_dir), reverse=True):
            if not filename.endswith(".json") or filename == "_last.json":
                continue
            path = os.path.join(self._data_dir, filename)
            data = safe_load_json(path, default=_CORRUPT)
            if data is _CORRUPT:
                continue
            convos.append(
                {
                    "id": data.get("id", filename[:-5]),
                    "title": data.get("title", ""),
                    "created_at": data.get("created_at", ""),
                    "updated_at": data.get("updated_at", ""),
                    "turns": data.get("total_messages_ever", 0) // 2,
                },
            )
            if len(convos) >= limit:
                break
        return convos

    def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a saved conversation from disk."""
        if not self._is_safe_id(conversation_id):
            return False
        path = os.path.join(self._data_dir, f"{conversation_id}.json")
        if not os.path.exists(path):
            return False
        try:
            os.remove(path)
            if self._get_last_id() == conversation_id:
                last_path = os.path.join(self._data_dir, "_last.json")
                with contextlib.suppress(OSError):
                    os.remove(last_path)
            with self._lock:
                if self._conversation_id == conversation_id:
                    self.messages = []
                    self.summary = ""
                    self._conversation_id = None
            return True
        except OSError as e:
            log.warning("Failed to delete conversation %s: %s", conversation_id, e)
            return False

    def rename_conversation(self, conversation_id: str, new_title: str) -> bool:
        """Rename a saved conversation. Returns True on success."""
        if not self._is_safe_id(conversation_id):
            return False
        path = os.path.join(self._data_dir, f"{conversation_id}.json")
        data = safe_load_json(path, default=_CORRUPT)
        if data is _CORRUPT:
            return False
        data["title"] = new_title[:200]
        try:
            atomic_write_json(path, data)
            with self._lock:
                if self._conversation_id == conversation_id:
                    self._title = new_title[:200]
            return True
        except OSError as e:
            log.warning("Failed to rename conversation %s: %s", conversation_id, e)
            return False

    def search_conversations(self, query: str, limit: int = 20) -> list[dict]:
        """Search across all saved conversations for a query string."""
        os.makedirs(self._data_dir, exist_ok=True)
        q = query.lower()
        results = []
        for filename in sorted(os.listdir(self._data_dir), reverse=True):
            if not filename.endswith(".json") or filename == "_last.json":
                continue
            path = os.path.join(self._data_dir, filename)
            data = safe_load_json(path, default=_CORRUPT)
            if data is _CORRUPT:
                continue
            title = data.get("title", "")
            messages = data.get("messages", [])
            # Search in title and message contents
            found = q in (title or "").lower()
            if not found:
                for m in messages:
                    if q in m.get("content", "").lower():
                        found = True
                        break
            if found:
                results.append(
                    {
                        "id": data.get("id", filename[:-5]),
                        "title": title,
                        "created_at": data.get("created_at", ""),
                        "updated_at": data.get("updated_at", ""),
                        "turns": data.get("total_messages_ever", 0) // 2,
                    },
                )
                if len(results) >= limit:
                    break
        return results

    def export_markdown(self) -> str:
        """Export current conversation as Markdown text."""
        with self._lock:
            lines = [f"# {self._title or 'Conversation'}"]
            if self._created_at:
                lines.append(f"\n*Started: {self._created_at}*\n")
            if self.summary:
                lines.append("## Summary of earlier messages\n")
                lines.append(self.summary)
                lines.append("")
            lines.append("---\n")
            for m in self.messages:
                role = "**User:**" if m["role"] == "user" else "**Assistant:**"
                lines.append(f"{role} {m['content']}\n")
            return "\n".join(lines)

    def export_all_markdown(self) -> str:
        """Export all saved conversations as a single Markdown document."""
        os.makedirs(self._data_dir, exist_ok=True)
        parts = ["# All Conversations\n"]
        for filename in sorted(os.listdir(self._data_dir)):
            if not filename.endswith(".json") or filename == "_last.json":
                continue
            path = os.path.join(self._data_dir, filename)
            data = safe_load_json(path, default=_CORRUPT)
            if data is _CORRUPT:
                continue
            title = data.get("title", filename[:-5])
            parts.append(f"## {title}")
            created = data.get("created_at", "")
            if created:
                parts.append(f"*{created}*\n")
            for m in data.get("messages", []):
                role = "**User:**" if m["role"] == "user" else "**Assistant:**"
                parts.append(f"{role} {m['content']}\n")
            parts.append("---\n")
        return "\n".join(parts)

    @property
    def total_turns_ever(self) -> int:
        """Total number of turns across all messages (user+assistant pairs)."""
        with self._lock:
            return self._total_messages_ever // 2

    def _save_last_id(self) -> None:
        path = os.path.join(self._data_dir, "_last.json")
        with contextlib.suppress(OSError):
            atomic_write_json(path, {"id": self._conversation_id})

    def _get_last_id(self) -> str | None:
        path = os.path.join(self._data_dir, "_last.json")
        data = safe_load_json(path, default={})
        conv_id = data.get("id")
        return str(conv_id) if conv_id is not None else None
