import json
import time
import os
import re
import shutil
import asyncio
from pathlib import Path
from typing import Dict, Any, Optional
from loguru import logger

from core.paths import PERSONA_DIR

SESSION_DIR = PERSONA_DIR / "sessions"
SESSION_FILE = SESSION_DIR / "sessions.json"
LOGS_DIR = SESSION_DIR / "logs"
EVENTS_DIR = SESSION_DIR / "events"
SKILLS_DIR = Path("skills")
_STRUCTURAL_RESIDUE_RE = re.compile(r"^[`{}\[\],:;\s]+$")
_LEGACY_XML_TOOL_TAG_NAMES = (
    "read_file|write_file|delete_file|list_dir|search_files|run_command|"
    "memory_search|google_search|browser_navigate|spawn_agent|"
    "save_memory|log_memory|ls|dir|list_files|cat|open_file|show_file|"
    "grep|rg|ripgrep|find_files|shell|terminal|exec|bash|powershell|cmd|tool_code"
)
_LEGACY_XML_TOOL_OPEN_RE = re.compile(
    rf"<(?:{_LEGACY_XML_TOOL_TAG_NAMES})>",
    re.IGNORECASE,
)
_LEGACY_XML_TOOL_BLOCK_RE = re.compile(
    rf"<(?:{_LEGACY_XML_TOOL_TAG_NAMES})>\s*.*?\s*</(?:{_LEGACY_XML_TOOL_TAG_NAMES})>",
    re.IGNORECASE | re.DOTALL,
)
_PIPE_TOOL_OPEN_RE = re.compile(
    rf"<\|(?:{_LEGACY_XML_TOOL_TAG_NAMES})\|>",
    re.IGNORECASE,
)
_PIPE_TOOL_BLOCK_RE = re.compile(
    rf"<\|(?:{_LEGACY_XML_TOOL_TAG_NAMES})\|>\s*.*?\s*<\|/(?:{_LEGACY_XML_TOOL_TAG_NAMES})\|>",
    re.IGNORECASE | re.DOTALL,
)
_INLINE_TOOL_CALL_RE = re.compile(
    r"(?::?functions\.\w+(?::\d+)?\s*\{[^}]*\})",
    re.IGNORECASE,
)
_JSON_TOOL_OBJECT_RE = re.compile(
    r'\{\s*"name"\s*:\s*"[^"]+"\s*,\s*"(?:arguments|parameters)"\s*:\s*\{[^}]*\}\s*\}',
    re.IGNORECASE | re.DOTALL,
)


def _sanitize_assistant_text_content(content: str) -> tuple[str, bool]:
    """Remove assistant-side leaked tool syntax from persisted history."""
    if not isinstance(content, str) or not content:
        return content, False

    cleaned = content
    marker_positions = []

    for pattern in (
        _LEGACY_XML_TOOL_OPEN_RE,
        _PIPE_TOOL_OPEN_RE,
        _INLINE_TOOL_CALL_RE,
        _JSON_TOOL_OBJECT_RE,
    ):
        match = pattern.search(cleaned)
        if match:
            marker_positions.append(match.start())

    if marker_positions:
        cleaned = cleaned[: min(marker_positions)]

    cleaned = _LEGACY_XML_TOOL_BLOCK_RE.sub("", cleaned)
    cleaned = _PIPE_TOOL_BLOCK_RE.sub("", cleaned)
    cleaned = _INLINE_TOOL_CALL_RE.sub("", cleaned)
    cleaned = _JSON_TOOL_OBJECT_RE.sub("", cleaned)
    cleaned = cleaned.strip()

    if _STRUCTURAL_RESIDUE_RE.fullmatch(cleaned or ""):
        cleaned = ""

    return cleaned, cleaned != content.strip()


def _sanitize_assistant_content(message: Dict[str, Any]) -> tuple[Optional[Dict[str, Any]], bool]:
    """Drop malformed assistant residue while preserving valid tool-call entries."""
    if message.get("role") != "assistant":
        return message, False

    content = message.get("content")
    if not isinstance(content, str):
        return message, False

    cleaned = dict(message)
    stripped, content_changed = _sanitize_assistant_text_content(content)

    if stripped:
        if content_changed:
            cleaned["content"] = stripped
            return cleaned, True
        return message, False

    if not content.strip() or _STRUCTURAL_RESIDUE_RE.fullmatch(content.strip()) or content_changed:
        cleaned["content"] = ""
        if cleaned.get("tool_calls"):
            return cleaned, True
        return None, True

    return message, False


def _sanitize_history_rows(rows: list[Dict[str, Any]]) -> tuple[list[Dict[str, Any]], int]:
    cleaned_rows: list[Dict[str, Any]] = []
    cleaned_count = 0
    for row in rows or []:
        cleaned_row, changed = _sanitize_assistant_content(row)
        if changed:
            cleaned_count += 1
        if cleaned_row is not None:
            cleaned_rows.append(cleaned_row)
    return cleaned_rows, cleaned_count


class SessionManager:
    def __init__(self):

        SESSION_DIR.mkdir(parents=True, exist_ok=True)
        LOGS_DIR.mkdir(parents=True, exist_ok=True)
        EVENTS_DIR.mkdir(parents=True, exist_ok=True)

        self.sessions: Dict[str, Any] = {}
        self._skills_cache: Optional[list] = None
        self._last_skills_refresh: float = 0
        self._lock = asyncio.Lock()
        self._save_task: Optional[asyncio.Task] = None
        self._save_debounce_s: float = 2.0

        self.sessions = self._load_sessions()

    def _load_sessions(self) -> Dict[str, Any]:
        """Load sessions from JSON file."""
        if SESSION_FILE.exists():
            try:
                return json.loads(SESSION_FILE.read_text(encoding="utf-8"))
            except Exception as e:
                logger.error(f"Error loading sessions: {e}")
                return {}
        return {}

    async def _save_sessions(self):
        """Persist sessions to JSON file using an atomic write pattern."""
        async with self._lock:
            try:
                temp_file = SESSION_FILE.with_suffix(".tmp")
                content = await asyncio.to_thread(json.dumps, self.sessions)
                await asyncio.to_thread(temp_file.write_text, content, encoding="utf-8")
                await asyncio.to_thread(shutil.move, str(temp_file), str(SESSION_FILE))
            except Exception as e:
                logger.error(f"Error saving sessions: {e}")

    async def _save_sessions_debounced(self):
        if self._save_task and not self._save_task.done():
            return

        async def _delayed_save():
            await asyncio.sleep(self._save_debounce_s)
            await self._save_sessions()

        self._save_task = asyncio.create_task(_delayed_save())

    async def update_session(
        self,
        session_key: str,
        model: str,
        origin: str,
        usage: Optional[Any] = None,
        injected_files: Optional[list] = None,
        parent_id: Optional[str] = None,
        task: Optional[str] = None,
        subagent_name: Optional[str] = None,
    ):

        current_time = time.time()

        if session_key not in self.sessions:
            self.sessions[session_key] = {
                "id": session_key,
                "created_at": current_time,
                "last_active": current_time,
                "origin": origin,
                "model": model,
                "total_tokens": {"input": 0, "output": 0, "total": 0},
                "skills": self._get_skills(),
                "injected_files": injected_files or [],
                "history_file": f"persona/sessions/logs/{session_key}.jsonl",
                "parent_id": parent_id,
                "task": task,
                "subagent_name": subagent_name,
            }

        session = self.sessions[session_key]
        session["last_active"] = current_time
        session["model"] = model

        if parent_id:
            session["parent_id"] = parent_id
        if task:
            session["task"] = task
        if subagent_name:
            session["subagent_name"] = subagent_name

        if injected_files:
            existing = set(session.get("injected_files", []))
            for f in injected_files:
                existing.add(f)
            session["injected_files"] = sorted(list(existing))

        if usage:
            p_tokens = 0
            c_tokens = 0
            t_tokens = 0

            if isinstance(usage, dict):
                p_tokens = usage.get("prompt_tokens", 0)
                c_tokens = usage.get("completion_tokens", 0)
                t_tokens = usage.get("total_tokens", 0)
            else:
                p_tokens = getattr(usage, "prompt_tokens", 0)
                c_tokens = getattr(usage, "completion_tokens", 0)
                t_tokens = getattr(usage, "total_tokens", 0)

            if "total_tokens" not in session:
                session["total_tokens"] = {"input": 0, "output": 0, "total": 0}

            session["total_tokens"]["input"] += p_tokens
            session["total_tokens"]["output"] += c_tokens
            session["total_tokens"]["total"] += t_tokens

        await self._save_sessions_debounced()

    async def append_chat_log(self, session_key: str, message: Dict[str, Any]):
        """Append a message to the session's JSONL log (background)."""

        def _sync_append():
            log_file = LOGS_DIR / f"{session_key}.jsonl"
            try:
                with open(log_file, "a", encoding="utf-8") as f:
                    msg_with_time = message.copy()
                    msg_with_time["timestamp"] = time.time()
                    f.write(json.dumps(msg_with_time) + "\n")
            except Exception as e:
                logger.error(f"Error appending chat log: {e}")

        await asyncio.to_thread(_sync_append)

    async def append_event_log(self, session_key: str, event: Dict[str, Any]):
        """Append a structured event to the session audit log."""

        def _sync_append():
            log_file = EVENTS_DIR / f"{session_key}.jsonl"
            try:
                payload = dict(event)
                payload.setdefault("timestamp", time.time())
                with open(log_file, "a", encoding="utf-8") as f:
                    f.write(json.dumps(payload, default=str) + "\n")
            except Exception as e:
                logger.error(f"Error appending event log: {e}")

        await asyncio.to_thread(_sync_append)

    async def save_history(self, session_key: str, history: list):
        """Persist full conversation history to disk (JSON)."""
        history_dir = SESSION_DIR / "history"
        history_dir.mkdir(parents=True, exist_ok=True)
        history_file = history_dir / f"{session_key}.json"
        try:
            cleaned_history, _ = _sanitize_history_rows(history)
            content = await asyncio.to_thread(json.dumps, cleaned_history, default=str)
            await asyncio.to_thread(history_file.write_text, content, encoding="utf-8")
        except Exception as e:
            logger.error(f"Error saving history for {session_key}: {e}")

    async def load_history(self, session_key: str) -> list:
        """Load conversation history from disk. Returns [] if not found."""
        history_file = SESSION_DIR / "history" / f"{session_key}.json"
        if history_file.exists():
            try:
                content = await asyncio.to_thread(
                    history_file.read_text, encoding="utf-8"
                )
                history = json.loads(content)
                cleaned_history, cleaned_count = _sanitize_history_rows(history)
                if cleaned_count:
                    logger.info(
                        f"Cleaned {cleaned_count} malformed history entr{'y' if cleaned_count == 1 else 'ies'} in {session_key}."
                    )
                    await self.save_history(session_key, cleaned_history)
                return cleaned_history
            except Exception as e:
                logger.debug(f"Could not load history for {session_key}: {e}")
        return []

    def cleanup_history_artifacts(self) -> dict[str, int]:
        """Sanitize persisted histories and chat logs to remove malformed assistant residue."""
        history_dir = SESSION_DIR / "history"
        cleaned_history_files = 0
        cleaned_history_entries = 0
        cleaned_log_files = 0
        cleaned_log_entries = 0

        if history_dir.exists():
            for history_file in history_dir.glob("*.json"):
                try:
                    history = json.loads(history_file.read_text(encoding="utf-8"))
                except Exception:
                    continue
                cleaned_history, cleaned_count = _sanitize_history_rows(history)
                if not cleaned_count:
                    continue
                history_file.write_text(
                    json.dumps(cleaned_history, default=str), encoding="utf-8"
                )
                cleaned_history_files += 1
                cleaned_history_entries += cleaned_count

        if LOGS_DIR.exists():
            for log_file in LOGS_DIR.glob("*.jsonl"):
                rows = []
                changed = 0
                try:
                    with log_file.open("r", encoding="utf-8") as f:
                        for line in f:
                            line = line.strip()
                            if not line:
                                continue
                            try:
                                row = json.loads(line)
                            except Exception:
                                rows.append(line)
                                continue
                            cleaned_row, row_changed = _sanitize_assistant_content(row)
                            if row_changed:
                                changed += 1
                            if cleaned_row is not None:
                                rows.append(json.dumps(cleaned_row, default=str))
                except Exception:
                    continue
                if not changed:
                    continue
                log_file.write_text("\n".join(rows) + "\n", encoding="utf-8")
                cleaned_log_files += 1
                cleaned_log_entries += changed

        return {
            "history_files": cleaned_history_files,
            "history_entries": cleaned_history_entries,
            "log_files": cleaned_log_files,
            "log_entries": cleaned_log_entries,
        }

    def delete_history(self, session_key: str):
        """Remove persisted history file for a session."""
        history_file = SESSION_DIR / "history" / f"{session_key}.json"
        if history_file.exists():
            try:
                os.remove(history_file)
            except Exception:
                pass

    async def delete_session(self, session_key: str) -> bool:
        """Delete a session and its logs."""
        async with self._lock:
            if session_key in self.sessions:
                del self.sessions[session_key]

                temp_file = SESSION_FILE.with_suffix(".tmp")
                content = await asyncio.to_thread(json.dumps, self.sessions)
                await asyncio.to_thread(temp_file.write_text, content, encoding="utf-8")
                await asyncio.to_thread(shutil.move, str(temp_file), str(SESSION_FILE))

                log_file = LOGS_DIR / f"{session_key}.jsonl"
                if log_file.exists():
                    try:
                        await asyncio.to_thread(os.remove, log_file)
                    except Exception as e:
                        logger.error(f"Error deleting log file {log_file}: {e}")

                event_file = EVENTS_DIR / f"{session_key}.jsonl"
                if event_file.exists():
                    try:
                        await asyncio.to_thread(os.remove, event_file)
                    except Exception as e:
                        logger.error(f"Error deleting event file {event_file}: {e}")

                self.delete_history(session_key)
                return True
        return False

    async def delete_sessions(self, session_keys: list[str]) -> int:
        """Delete multiple sessions and their logs efficiently."""
        deleted_count = 0
        async with self._lock:
            for key in session_keys:
                if key in self.sessions:
                    del self.sessions[key]
                    deleted_count += 1

                    # Delete logs
                    log_file = LOGS_DIR / f"{key}.jsonl"
                    if log_file.exists():
                        try:
                            await asyncio.to_thread(os.remove, log_file)
                        except Exception as e:
                            logger.error(f"Error deleting log file {log_file}: {e}")

                    event_file = EVENTS_DIR / f"{key}.jsonl"
                    if event_file.exists():
                        try:
                            await asyncio.to_thread(os.remove, event_file)
                        except Exception as e:
                            logger.error(f"Error deleting event file {event_file}: {e}")

                    # Delete history
                    self.delete_history(key)

            if deleted_count > 0:
                # Save once after all deletions
                temp_file = SESSION_FILE.with_suffix(".tmp")
                content = await asyncio.to_thread(json.dumps, self.sessions)
                await asyncio.to_thread(temp_file.write_text, content, encoding="utf-8")
                await asyncio.to_thread(shutil.move, str(temp_file), str(SESSION_FILE))

        return deleted_count

    def get_sessions(self) -> Dict[str, Any]:
        """Return all sessions from memory snapshot."""
        current_skills = self._get_skills()
        for session in self.sessions.values():
            session["skills"] = current_skills
        return self.sessions

    def _get_skills(self) -> list:
        """Scan skills directory for available skills (cached 30s)."""
        now = time.time()
        if self._skills_cache is not None and now - self._last_skills_refresh < 30:
            return self._skills_cache

        skills = []
        enabled = []
        try:
            config_path = Path("limebot.json")
            if config_path.exists():
                data = json.loads(config_path.read_text(encoding="utf-8"))
                enabled = data.get("skills", {}).get("enabled", [])
        except Exception:
            pass

        if SKILLS_DIR.exists():
            for item in SKILLS_DIR.iterdir():
                if item.is_dir() and not item.name.startswith("__"):
                    if item.name in enabled:
                        skills.append(item.name)

        self._skills_cache = sorted(skills)
        self._last_skills_refresh = now
        return self._skills_cache
