"""Session management — Session class, get/rollback/branch/save/restore.

Extracted from core.py to reduce god object. All core.py dependencies
use lazy imports to avoid circular import issues.
"""

import json
import re
import threading
import time
import uuid as _uuid
from datetime import datetime
from typing import Optional

from salmalm.constants import DATA_DIR, KST
from salmalm.security.crypto import log


# ── Lazy accessors for core.py globals (break circular import) ──


def _get_db():
    """Get db."""
    from salmalm.core.core import _get_db as _impl

    return _impl()


def _audit_log(*args, **kwargs):
    """Audit log."""
    from salmalm.core.core import audit_log as _impl

    return _impl(*args, **kwargs)


def _restore_compaction_summary(session_id: str) -> Optional[str]:
    """Restore compaction summary."""
    from salmalm.core.compaction import _restore_compaction_summary as _impl

    return _impl(session_id)


_PERSIST_MESSAGE_LIMIT = 50  # Max messages to persist to SQLite


class Session:
    """OpenClaw-style isolated session with its own context.

    Each session has:
    - Unique ID and isolated message history
    - No cross-contamination between sessions
    - Automatic memory flush before compaction
    - Session metadata tracking
    """

    def __init__(self, session_id: str, user_id: Optional[int] = None) -> None:
        """Create a new session with isolated context, message history, and model preferences."""
        self.id = session_id
        self.user_id = user_id  # Multi-tenant: owning user (None = legacy/local)
        self.messages: list = []
        self.created = time.time()
        self.last_active = time.time()
        self.metadata: dict = {}  # Arbitrary session metadata
        self._memory_flushed = False  # Track if pre-compaction memory flush happened
        self.thinking_enabled = False  # Extended thinking toggle (default OFF)
        self.thinking_level = "medium"  # Thinking depth: "low"|"medium"|"high"|"xhigh"
        self.model_override = None  # None=inherit force_model; 'auto'=explicit auto routing; or specific model
        self.tts_enabled = False  # TTS toggle
        self.tts_voice = "alloy"  # TTS voice selection
        self._thinking_suggested = False  # Track if thinking suggestion shown this session
        self.last_model = "auto"  # Last used model (for UI display)
        self.last_complexity = "auto"  # Last complexity level

    def add_system(self, content: str) -> None:
        # Replace existing system message
        """Add a system message to the session."""
        _MAX_SYSTEM_CHARS = 200_000  # 200KB hard cap — generous for complex system prompts
        if len(content) > _MAX_SYSTEM_CHARS:
            content = content[:_MAX_SYSTEM_CHARS]
            log.warning("[SESSION] System prompt truncated to %d chars for %s", _MAX_SYSTEM_CHARS, self.id)
        self.messages = [m for m in self.messages if m["role"] != "system"]
        self.messages.insert(0, {"role": "system", "content": content})

    def _persist(self):
        """Save session to SQLite (only text messages, skip image data).

        Handles: disk full (OSError), DB lock (sqlite3.OperationalError).
        """
        try:
            # Filter out large binary data from messages
            saveable = []
            for m in self.messages[-_PERSIST_MESSAGE_LIMIT:]:  # Keep last 50 messages
                if isinstance(m.get("content"), list):
                    # Multimodal — replace base64 images with placeholder, keep text
                    parts = []
                    for b in m["content"]:
                        if b.get("type") == "text":
                            parts.append(b)
                        elif b.get("type") == "image":
                            # Replace heavy base64 with marker so context is preserved
                            parts.append({"type": "text", "text": "[Image was attached to this message]"})
                    if parts:
                        saveable.append({**m, "content": parts})
                elif isinstance(m.get("content"), str):
                    saveable.append(m)
            conn = _get_db()
            # Persist session metadata (model_override, thinking, tts)
            _meta = json.dumps(
                {
                    "model_override": self.model_override if self.model_override and self.model_override != "auto" else None,
                    "thinking_enabled": self.thinking_enabled,
                    "thinking_level": self.thinking_level,
                    "tts_enabled": self.tts_enabled,
                    "tts_voice": self.tts_voice,
                },
                ensure_ascii=False,
            )
            # Use UPSERT (INSERT ... ON CONFLICT DO UPDATE) instead of INSERT OR REPLACE.
            # INSERT OR REPLACE = DELETE + INSERT → wipes columns not in the INSERT list
            # (title, parent_session_id, branch_index → all reset to DEFAULT on every save).
            # UPSERT only overwrites the columns we explicitly set, preserving the rest.
            conn.execute(
                """INSERT INTO session_store (session_id, messages, updated_at, user_id, session_meta)
                   VALUES (?,?,?,?,?)
                   ON CONFLICT(session_id) DO UPDATE SET
                       messages=excluded.messages,
                       updated_at=excluded.updated_at,
                       user_id=excluded.user_id,
                       session_meta=excluded.session_meta""",
                (
                    self.id,
                    json.dumps(saveable, ensure_ascii=False),
                    datetime.now(KST).isoformat(),
                    self.user_id,
                    _meta,
                ),
            )  # noqa: F405
            conn.commit()
        except Exception as e:
            log.warning(f"Session persist error: {e}")

    def add_user(self, content: str) -> None:
        """Add a user message to the session.

        Auto-compaction: if session exceeds 1000 messages, trim old ones.
        """
        self.messages.append({"role": "user", "content": content})
        self.last_active = time.time()
        # Session size explosion prevention
        if len(self.messages) > 1000:
            system_msgs = [m for m in self.messages if m["role"] == "system"][:1]
            recent = [m for m in self.messages if m["role"] != "system"][-50:]
            self.messages = system_msgs + recent
            log.warning(f"[SESSION] Auto-trimmed session {self.id}: >1000 msgs → {len(self.messages)}")
        # Persist user message immediately — prevents message loss on crash
        # before add_assistant() is called (e.g. long LLM round-trip with kill -9)
        try:
            self._persist()
        except Exception as _e:
            log.warning("[SESSION] add_user _persist failed for %s: %s", self.id, _e)

    def add_assistant(self, content: str) -> None:
        """Add an assistant response to the session."""
        self.messages.append({"role": "assistant", "content": content})
        self.last_active = time.time()
        self._persist()
        # Auto-save to disk after final response (debounced — not on tool calls)
        try:
            save_session_to_disk(self.id)
        except Exception as e:
            log.debug(f"Suppressed: {e}")

    def add_tool_results(self, results: list) -> None:
        """Add tool results as a single user message with all results.
        results: list of {'tool_use_id': str, 'content': str}
        """
        content = [
            {
                "type": "tool_result",
                "tool_use_id": r["tool_use_id"],
                "content": r["content"],
            }
            for r in results
        ]
        self.messages.append({"role": "user", "content": content})


_tg_bot = None  # Set during startup by telegram module


def get_telegram_bot() -> Optional[object]:
    """Accessor for the Telegram bot instance (avoids direct global access)."""
    return _tg_bot


def set_telegram_bot(bot: object) -> None:
    """Set the Telegram bot instance (called during startup)."""
    global _tg_bot
    _tg_bot = bot


_llm_cron = None  # Set during startup by __main__ (LLMCron instance)
_sessions = {}  # type: ignore[var-annotated]
_session_lock = threading.Lock()  # Protects _sessions dict
_session_cleanup_ts = 0.0
_SESSION_TTL = 3600 * 8  # 8 hours
_SESSION_MAX = 200


def _cleanup_sessions():
    """Remove inactive sessions older than TTL."""
    global _session_cleanup_ts
    now = time.time()
    if now - _session_cleanup_ts < 600:  # Check every 10 min
        return
    _session_cleanup_ts = now
    with _session_lock:
        stale = [sid for sid, s in _sessions.items() if now - s.last_active > _SESSION_TTL]
        for sid in stale:
            try:
                _sessions[sid]._persist()
            except Exception as e:
                log.debug(f"Suppressed: {e}")
            del _sessions[sid]
            try:
                from salmalm.core.session_manager import evict_session_timing
                evict_session_timing(sid)
            except Exception:
                pass
        if stale:
            log.info(f"[CLEAN] Session cleanup: removed {len(stale)} inactive sessions")
        # Hard cap
        if len(_sessions) > _SESSION_MAX:
            by_age = sorted(_sessions.items(), key=lambda x: x[1].last_active)
            for sid, _ in by_age[: len(_sessions) - _SESSION_MAX]:
                del _sessions[sid]
                try:
                    from salmalm.core.session_manager import evict_session_timing
                    evict_session_timing(sid)
                except Exception:
                    pass


def get_session(session_id: str, user_id: Optional[int] = None) -> Session:
    """Get or create a chat session by ID.

    If user_id is provided and the session already exists with a different user_id,
    access is denied (returns a new isolated session instead of the existing one).
    """
    # Sanitize session_id: prevent path traversal and memory key injection
    if not session_id or len(session_id) > 128:
        session_id = "default"
    elif re.search(r'[/\\\x00\r\n]|\.\.', session_id):
        session_id = re.sub(r"[^a-zA-Z0-9_\-\.@:]", "_", session_id)[:128]
    _cleanup_sessions()
    with _session_lock:
        if session_id in _sessions:
            existing = _sessions[session_id]
            # Access control: deny if owned by a *different* authenticated user.
            if user_id is not None:
                _owner_conflict = (
                    existing.user_id is not None and existing.user_id != user_id
                ) or (
                    existing.user_id is None
                    and session_id not in ("web", "local", "default")
                    and _is_multi_user_mode()
                )
                if _owner_conflict:
                    log.warning(
                        "[SESSION] User %s denied access to session %s (owned by %s)",
                        user_id, session_id, existing.user_id
                    )
                    isolated_id = f"{session_id}_u{user_id}"
                    if isolated_id not in _sessions:
                        _sessions[isolated_id] = Session(isolated_id, user_id=user_id)
                    return _sessions[isolated_id]
            return existing
        if session_id not in _sessions:
            _sessions[session_id] = Session(session_id, user_id=user_id)
            # Try to restore from SQLite
            try:
                conn = _get_db()
                row = conn.execute(
                    "SELECT messages, session_meta FROM session_store WHERE session_id=?",
                    (session_id,),
                ).fetchone()
                if row:
                    try:
                        restored = json.loads(row[0])
                        if not isinstance(restored, list):
                            raise ValueError("Session data is not a list")
                        # Sanitize: only allow known roles from DB
                        _ALLOWED_ROLES = {"user", "assistant", "system", "tool"}
                        sanitized = [m for m in restored if isinstance(m, dict) and m.get("role") in _ALLOWED_ROLES]
                        if len(sanitized) != len(restored):
                            log.warning("[SESSION] Dropped %d invalid-role messages from DB for %s",
                                        len(restored) - len(sanitized), session_id)
                        _sessions[session_id].messages = sanitized
                        log.info(f"[NOTE] Session restored: {session_id} ({len(sanitized)} msgs)")
                    except (json.JSONDecodeError, ValueError, TypeError) as je:
                        # Corrupted session JSON — start fresh
                        log.warning(f"[SESSION] Corrupt session JSON for {session_id}: {je}")
                        _sessions[session_id].messages = []
                    # Restore session metadata (model_override, thinking, tts)
                    try:
                        _meta_json = row[1] if len(row) > 1 and row[1] else "{}"
                        _meta = json.loads(_meta_json)
                        _saved_ov = _meta.get("model_override")
                        if _saved_ov and _saved_ov != "auto":
                            _sessions[session_id].model_override = _saved_ov
                        if _meta.get("thinking_enabled") is not None:
                            _sessions[session_id].thinking_enabled = _meta["thinking_enabled"]
                        if _meta.get("thinking_level"):
                            _sessions[session_id].thinking_level = _meta["thinking_level"]
                        if _meta.get("tts_enabled") is not None:
                            _sessions[session_id].tts_enabled = _meta["tts_enabled"]
                        if _meta.get("tts_voice"):
                            _sessions[session_id].tts_voice = _meta["tts_voice"]
                    except (json.JSONDecodeError, IndexError, TypeError) as me:
                        log.debug(f"Session meta restore: {me}")
                    # Refresh system prompt
                    from salmalm.core.prompt import build_system_prompt

                    _sessions[session_id].add_system(build_system_prompt(full=False))
                    return _sessions[session_id]
            except Exception as e:
                log.warning(f"Session restore error: {e}")
            from salmalm.core.prompt import build_system_prompt

            _sessions[session_id].add_system(build_system_prompt(full=True))

            # Cross-session continuity: inject last compaction summary
            prev_summary = _restore_compaction_summary(session_id)
            if prev_summary:
                _sessions[session_id].messages.append(
                    {
                        "role": "system",
                        "content": f"[Previous session context]\n{prev_summary}",
                    }
                )
                log.info(f"[NOTE] Restored compaction summary for {session_id} ({len(prev_summary)} chars)")

            # Apply onboarding model as session default
            try:
                from salmalm.security.crypto import vault

                if vault.is_unlocked:
                    dm = vault.get("default_model")
                    if dm and dm != "auto":
                        _sessions[session_id]._default_model = dm
                        # Don't set model_override here — let user's explicit
                        # UI selection (persisted in DB) take precedence.
                        # _default_model is used as fallback in auto-routing only.
            except Exception as e:
                log.debug(f"Suppressed: {e}")

            log.info(
                f"[NOTE] New session: {session_id} (system prompt: {len(_sessions[session_id].messages[0]['content'])} chars)"
            )
            _audit_log(
                "session_create",
                f"new session: {session_id}",
                session_id=session_id,
                detail_dict={"session_id": session_id},
            )
        try:
            from salmalm.monitoring.metrics import active_sessions
            active_sessions.set(len(_sessions))
        except Exception:
            pass
        return _sessions[session_id]


def rollback_session(session_id: str, count: int) -> dict:
    """Roll back the last `count` user+assistant message pairs.

    Removed messages are backed up in session_message_backup table.
    Returns {'ok': True, 'removed': <int>} or {'ok': False, 'error': ...}.
    """
    session = get_session(session_id)
    non_system = [(i, m) for i, m in enumerate(session.messages) if m.get("role") != "system"]
    pairs_removed = 0
    indices_to_remove = []
    j = len(non_system) - 1
    while pairs_removed < count and j >= 0:
        idx_j, msg_j = non_system[j]
        if msg_j["role"] == "assistant":
            indices_to_remove.append(idx_j)
            if j - 1 >= 0 and non_system[j - 1][1]["role"] == "user":
                indices_to_remove.append(non_system[j - 1][0])
                j -= 2
            else:
                j -= 1
            pairs_removed += 1
        elif msg_j["role"] == "user":
            indices_to_remove.append(idx_j)
            j -= 1
            pairs_removed += 1
        else:
            j -= 1

    if not indices_to_remove:
        return {"ok": False, "error": "No messages to rollback"}

    removed_msgs = [session.messages[i] for i in sorted(indices_to_remove)]
    conn = _get_db()
    conn.execute(
        "INSERT INTO session_message_backup (session_id, messages_json, removed_at, reason) VALUES (?,?,?,?)",
        (
            session_id,
            json.dumps(removed_msgs, ensure_ascii=False),
            datetime.now(KST).isoformat(),
            "rollback",
        ),
    )  # noqa: F405
    conn.commit()

    for i in sorted(indices_to_remove, reverse=True):
        session.messages.pop(i)

    session._persist()
    try:
        save_session_to_disk(session_id)
    except Exception as e:
        log.debug(f"Suppressed: {e}")
    _audit_log("session_rollback", f"{session_id}: removed {pairs_removed} pairs")
    return {"ok": True, "removed": pairs_removed}


def branch_session(session_id: str, message_index: int, user_id: int | None = None) -> dict:
    """Create a new session branching from session_id at message_index.

    Copies messages[0:message_index+1] into a new session.
    Returns {'ok': True, 'new_session_id': ...} or error.
    """

    session = get_session(session_id)
    # Ownership check: non-admin users can only branch their own sessions
    if user_id is not None and session.user_id is not None and session.user_id != user_id:
        log.warning("[BRANCH] User %s denied branch of session %s (owned by %s)", user_id, session_id, session.user_id)
        return {"ok": False, "error": "Access denied — session belongs to another user"}
    if message_index < 0 or message_index >= len(session.messages):
        return {"ok": False, "error": f"Invalid message_index: {message_index}"}

    new_id = f"branch-{_uuid.uuid4().hex[:8]}"
    new_session = Session(new_id)
    new_session.messages = json.loads(json.dumps(session.messages[: message_index + 1]))
    new_session.metadata["parent_session_id"] = session_id
    new_session.metadata["branch_index"] = message_index

    with _session_lock:
        _sessions[new_id] = new_session

    new_session._persist()

    conn = _get_db()
    conn.execute(
        "UPDATE session_store SET parent_session_id=?, branch_index=? WHERE session_id=?",
        (session_id, message_index, new_id),
    )
    conn.commit()

    try:
        save_session_to_disk(new_id)
    except Exception as e:
        log.debug(f"Suppressed: {e}")
    _audit_log("session_branch", f"{session_id} -> {new_id} at index {message_index}")
    return {"ok": True, "new_session_id": new_id, "parent_session_id": session_id}


_SESSIONS_DIR = DATA_DIR / "sessions"


def save_session_to_disk(session_id: str) -> None:
    """Serialize session state to ~/.salmalm/sessions/{id}.json."""
    with _session_lock:
        session = _sessions.get(session_id)
        if not session:
            return
    try:
        _SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
        saveable_msgs = []
        for m in session.messages[-_PERSIST_MESSAGE_LIMIT:]:
            if isinstance(m.get("content"), list):
                texts = [b for b in m["content"] if isinstance(b, dict) and b.get("type") == "text"]
                if texts:
                    saveable_msgs.append({**m, "content": texts})
            elif isinstance(m.get("content"), str):
                saveable_msgs.append(m)
        data = {
            "session_id": session.id,
            "messages": saveable_msgs,
            "created": session.created,
            "last_active": session.last_active,
            "metadata": session.metadata,
        }
        # Path traversal guard: sanitize session_id before using as filename
        _safe_sid = re.sub(r"[^a-zA-Z0-9_\-\.]", "_", session_id)[:128]
        path = _SESSIONS_DIR / f"{_safe_sid}.json"
        path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    except Exception as e:
        log.warning(f"[DISK] Failed to save session {session_id}: {e}")


def restore_session(session_id: str) -> Optional[Session]:
    """Load session from ~/.salmalm/sessions/{id}.json."""
    path = _SESSIONS_DIR / f"{session_id}.json"
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        session = Session(session_id)
        session.messages = data.get("messages", [])
        session.created = data.get("created", time.time())
        session.last_active = data.get("last_active", time.time())
        session.metadata = data.get("metadata", {})
        with _session_lock:
            _sessions[session_id] = session
        log.info(f"[DISK] Restored session from disk: {session_id} ({len(session.messages)} msgs)")
        return session
    except Exception as e:
        log.warning(f"[DISK] Failed to restore session {session_id}: {e}")
        return None


def restore_all_sessions_from_disk() -> None:
    """On startup, restore all sessions from disk."""
    if not _SESSIONS_DIR.exists():
        return
    count = 0
    for path in _SESSIONS_DIR.glob("*.json"):
        sid = path.stem
        if restore_session(sid):
            count += 1
    if count:
        log.info(f"[DISK] Restored {count} sessions from disk")
