## @file context.py
#  @brief Dual-stack conversation context (jsonl/md_as_config) with Markdown-first persistence.

from __future__ import annotations

import traceback
import datetime as _dt
from pathlib import Path
from typing import Iterable, List, Union, Optional, Callable, Any, Dict

from .model import Message, CompactRecord
from .session_store import SessionStore, DEFAULT_META_DIRNAME


class Context:
    @staticmethod
    def _store_options(*, meta_dirname: str, global_sessions_root: Path | None) -> Dict[str, Any]:
        return {
            "meta_dirname": meta_dirname,
            "global_sessions_root": global_sessions_root,
        }

    @classmethod
    def _from_store_path(
        cls,
        file_backend: Path,
        *,
        session_id: str | None = None,
        work_dir: Path | None = None,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
        restore: bool = False,
    ) -> "Context":
        ctx = cls(
            file_backend,
            session_id=session_id,
            work_dir=work_dir,
            **cls._store_options(meta_dirname=meta_dirname, global_sessions_root=global_sessions_root),
        )
        if restore:
            ctx.restore()
        return ctx

    def __init__(
        self,
        file_backend: Path,
        session_id: str | None = None,
        work_dir: Path | None = None,
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ):
        self._store = SessionStore(
            file_backend,
            session_id=session_id,
            work_dir=work_dir,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        self._history: List[Union[Message, CompactRecord]] = []
        self._mcp_client: Optional[Any] = None
        self._mcp_tool_schemas: Dict[str, dict] = {}
        self._mcp_tool_availability: Dict[str, bool] = {}
        self._mcp_ready: bool = False
        self._mcp_status_message: str = ""
        self._mcp_last_ping: Optional[_dt.datetime] = None

    # --------------------
    # MCP helpers
    # --------------------

    def set_mcp_client(self, client: Any) -> None:
        self._mcp_client = client

    def get_mcp_client(self) -> Optional[Any]:
        return self._mcp_client

    def mark_mcp_ready(self, ready: bool, message: str = "", last_ping: Optional[_dt.datetime] = None) -> None:
        self._mcp_ready = ready
        if message:
            self._mcp_status_message = message
        if last_ping:
            self._mcp_last_ping = last_ping

    def is_mcp_ready(self) -> bool:
        return self._mcp_ready

    def update_mcp_tool_state(self, schemas: Dict[str, dict], availability: Dict[str, bool]) -> None:
        self._mcp_tool_schemas = schemas
        self._mcp_tool_availability = availability

    def clear_mcp_tool_state(self) -> None:
        """Clear cached MCP tool schemas and availability."""
        self._mcp_tool_schemas = {}
        self._mcp_tool_availability = {}

    def get_mcp_tool_schemas(self) -> Dict[str, dict]:
        return dict(self._mcp_tool_schemas)

    def get_mcp_tool_availability(self) -> Dict[str, bool]:
        return dict(self._mcp_tool_availability)

    def get_mcp_tool_definition(self, name: str) -> Optional[dict]:
        """Return stored MCP tool definition by name."""
        return self._mcp_tool_schemas.get(name)

    def get_mcp_status_message(self) -> str:
        return str(self._mcp_status_message)

    def set_mcp_last_ping(self, when: _dt.datetime) -> None:
        self._mcp_last_ping = when

    def get_mcp_last_ping(self) -> Optional[_dt.datetime]:
        return self._mcp_last_ping

    @property
    def history(self) -> Iterable[Union[Message, CompactRecord]]:
        return list(self._history)



    # =============================================
    # Persistence and history restoration operations
    # =============================================

    def restore(self) -> bool:
        """Restore in-memory history from the configured backend file."""
        self._history.clear()
        try:
            self._history = self._store.restore()
            return bool(self._history)
        except Exception:
            traceback.print_exc()
            return False


    def set_tool_registry(self, _tool_registry: Any) -> None:
        # Kept for compatibility. Templates are now resolved from model objects directly.
        return None

    def select_template_for_item(self, item) -> str:
        """Compatibility shim. Dumps now resolve templates from objects directly."""
        template = getattr(item, "__chronml_template__", None)
        return template if isinstance(template, str) else ""

    @property
    def history_file(self) -> Path:
        return self._store.file_backend

    @property
    def session_id(self) -> str:
        return self._store.session_id or self._store.file_backend.stem

    @property
    def work_dir(self) -> Path | None:
        return self._store.work_dir

    @classmethod
    def load(
        cls,
        file_backend: Path,
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ) -> "Context":
        """Create a context and immediately restore records from backend."""
        return cls._from_store_path(
            file_backend,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
            restore=True,
        )

    def _clone_to_backend(
        self,
        target_backend: Path,
        history_items: List,
    ) -> "Context":
        clone = self._from_store_path(
            target_backend,
            meta_dirname=self._store.meta_dirname,
            global_sessions_root=self._store.global_sessions_root,
        )
        # Reuse provided history items directly.
        clone._history = history_items.copy()
        clone._store.dump_history(
            clone._history,
            target_backend,
            template="",
        )
        return clone

    def fork(
        self,
        target_backend: Path,
    ) -> "Context":
        """Clone current context history to a new backend file."""
        return self._clone_to_backend(target_backend, self._history).with_session(
            self.session_id,
            self.work_dir,
        )

    def truncate_and_fork(
        self,
        split_index: int,
        target_backend: Path,
    ) -> "Context":
        """Fork a truncated history up to ``split_index`` into a new backend."""
        split_index = max(0, min(len(self._history), split_index))

        # Keep records in [0, split_index).
        truncated_history = self._history[:split_index]

        return self._clone_to_backend(target_backend, truncated_history).with_session(
            self.session_id,
            self.work_dir,
        )

    def with_session(self, session_id: str | None, work_dir: Path | None) -> "Context":
        self._store.session_id = session_id
        self._store.work_dir = work_dir
        return self

    def adopt(self, other: "Context") -> "Context":
        self._store = other._store
        self._history = other._history
        return self

    # =====================================
    # Message append and compaction helpers
    # =====================================

    def append_message(self, message: Message | Iterable[Message]):
        msgs = list(message) if isinstance(message, Iterable) and not isinstance(message, Message) else [message]  # type: ignore[arg-type]
        sanitized: list[Message] = []
        for m in msgs:
            if not isinstance(m, Message):
                continue
            # Treat only None or empty string as empty content.
            # Empty dict/object content is still valid for tool workflows.
            has_content = m.content is not None and not (isinstance(m.content, str) and not m.content.strip())
            
            if not (has_content or (getattr(m, "tool_calls", None) is not None)):
                continue
            sanitized.append(m)
        if not sanitized:
            return
        for m in sanitized:
            self._history.append(m)
            # Persist incrementally.
            try:
                self._append_history_item_to_file(m)
            except:
                traceback.print_exc()
                raise

    def append_compact(self, summary: str, user_intent: str = "") -> None:
        """Append a compaction marker to history and persist it."""
        rec = CompactRecord(value=str(summary), user_intent=str(user_intent))
        self._history.append(rec)
        # Persist incrementally.
        self._append_history_item_to_file(rec)

    # ========================================
    # Import/export helpers for md_as_config
    # ========================================

    def export_md_as_config(
        self,
        md_path: Path,
        template: Optional[str] = None,
    ) -> None:
        """
        Export current history as md_as_config markdown using a template.
        """
        tmpl = template or ""
        self._store.export_md_as_config(self._history, md_path, template=tmpl)

    def import_md_as_config(self, md_path: Path) -> bool:
        """
        Import md_as_config markdown into backend and sync in-memory history.
        """
        records = self._store.import_md_as_config(md_path)
        if not records:
            return False
        self._sync_history_from_records(records)
        return True

    def _sync_history_from_records(self, records: Iterable[Any]) -> None:
        synced: list[Message | CompactRecord] = []
        for item in records:
            if isinstance(item, (Message, CompactRecord)):
                synced.append(item)
        self._history = synced

    def _append_history_item_to_file(self, item, template_func: Optional[Callable] = None) -> None:
        """Append one history item to backend incrementally."""
        self._store.append_item(item, template_func=template_func)

    @classmethod
    def create_session(
        cls,
        work_dir: Path,
        session_id: str | None = None,
        *,
        defer_persist: bool = False,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ) -> "Context":
        store = SessionStore.create_for_workdir(
            work_dir,
            session_id,
            defer_persist=defer_persist,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        ctx = cls._from_store_path(
            store.file_backend,
            session_id=store.session_id,
            work_dir=store.work_dir,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        ctx._store = store
        return ctx

    def ensure_persisted(self) -> None:
        self._store.ensure_persisted()

    async def aensure_persisted(self) -> None:
        await self._store.aensure_persisted()

    @classmethod
    def continue_session(
        cls,
        work_dir: Path,
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ) -> "Context" | None:
        picked = SessionStore.continue_in_workdir(work_dir.absolute(), meta_dirname=meta_dirname)
        if picked is None:
            return None
        latest_session_id, latest_path = picked
        if latest_path.exists():
            return cls._from_store_path(
                latest_path,
                session_id=latest_session_id,
                work_dir=work_dir,
                meta_dirname=meta_dirname,
                global_sessions_root=global_sessions_root,
                restore=True,
            )
        return None

    @classmethod
    def resume_session(
        cls,
        work_dir: Path,
        session_id: str,
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ) -> "Context" | None:
        work_dir = work_dir.absolute()
        history_file = SessionStore.resolve_session_path(work_dir, session_id, meta_dirname=meta_dirname)
        if history_file is None:
            return None
        return cls._from_store_path(
            history_file,
            session_id=session_id,
            work_dir=work_dir,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
            restore=True,
        )

    @classmethod
    def list_sessions(
        cls,
        work_dir: Path,
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
    ) -> tuple[list[tuple[str, Path, float]], str | None]:
        return SessionStore.list_sessions(work_dir, meta_dirname=meta_dirname)

    @classmethod
    def list_global_sessions(
        cls,
        work_dir: Path,
        *,
        global_sessions_root: Path | None = None,
        meta_dirname: str = DEFAULT_META_DIRNAME,
    ) -> tuple[list[tuple[str, Path, float]], str | None]:
        return SessionStore.list_global_sessions(
            work_dir,
            global_sessions_root=global_sessions_root,
            meta_dirname=meta_dirname,
        )

    @classmethod
    def fork_session(
        cls,
        work_dir: Path,
        source_session_id: str,
        new_session_id: str | None = "[AUTO]",
        *,
        meta_dirname: str = DEFAULT_META_DIRNAME,
        global_sessions_root: Path | None = None,
    ) -> "Context":
        work_dir = work_dir.absolute()
        source = cls.resume_session(
            work_dir,
            source_session_id,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        if source is None:
            raise ValueError(f"Session {source_session_id} not found to fork from")
        target_id = new_session_id
        if not target_id or target_id == "" or target_id == "[AUTO]":
            counter = 1
            while True:
                candidate = f"{source_session_id}-{counter}"
                if not SessionStore.session_exists(work_dir, candidate, meta_dirname=meta_dirname):
                    target_id = candidate
                    break
                counter += 1
        new_ctx = cls.create_session(
            work_dir,
            target_id,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        source.fork(new_ctx.history_file)
        new_ctx = cls.resume_session(
            work_dir,
            target_id,
            meta_dirname=meta_dirname,
            global_sessions_root=global_sessions_root,
        )
        return new_ctx
