# -*- coding: utf-8 -*-
"""Safe JSON session with filename sanitization for cross-platform
compatibility.

Windows filenames cannot contain: \\ / : * ? " < > |
"""
import os
import re
import logging
import shutil

from pathlib import Path
from typing import Union, Sequence

from qwenpaw.exceptions import ConfigurationException
from ...exceptions import AgentStateError
from ...utils.io_utils import (
    get_path_lock,
    run_sync_io,
    write_json_atomic_async,
)
from ...utils.json_utils import safe_json_loads as _safe_json_loads

logger = logging.getLogger(__name__)

# Characters forbidden in Windows filenames
_UNSAFE_FILENAME_RE = re.compile(r'[\\/:*?"<>|]')


def sanitize_filename(name: str) -> str:
    """Replace characters that are illegal in Windows filenames with ``--``.

    >>> sanitize_filename('discord:dm:12345')
    'discord--dm--12345'
    >>> sanitize_filename('normal-name')
    'normal-name'
    """
    return _UNSAFE_FILENAME_RE.sub("--", name)


# Marker used by ``sanitize_filename`` for the historical ``weixin:`` and
# canonical ``wechat:`` session_id prefixes.
_LEGACY_WEIXIN_SAFE_PREFIX = "weixin--"
_CANONICAL_WECHAT_SAFE_PREFIX = "wechat--"

# Sub-directory inside ``save_dir`` where the original legacy weixin
# session files are archived after migration. Keeping a copy preserves
# user data for manual recovery; the archive directory is excluded from
# regular session scans because callers list ``*.json`` non-recursively.
_WEIXIN_LEGACY_ARCHIVE_DIR = ".weixin-legacy"


def _read_session_json(path: str) -> dict:
    """Read and parse one session snapshot in a worker thread."""
    content = Path(path).read_text(
        encoding="utf-8",
        errors="surrogatepass",
    )
    return _safe_json_loads(content, path)


def migrate_legacy_weixin_session_files(save_dir: str) -> None:
    """Rename legacy ``weixin--`` session files to the ``wechat--`` form.

    Originals are moved to the ``.weixin-legacy/`` archive sub-dir so
    later startups skip the migration without extra bookkeeping. If a
    canonical file already exists, the legacy file is only archived.
    """
    if not save_dir or not os.path.isdir(save_dir):
        return
    try:
        entries = os.listdir(save_dir)
    except OSError:
        return
    legacy_files = [
        name
        for name in entries
        if name.endswith(".json")
        and _rewrite_weixin_in_session_filename(name) is not None
    ]
    if not legacy_files:
        return
    archive_dir = os.path.join(save_dir, _WEIXIN_LEGACY_ARCHIVE_DIR)
    try:
        os.makedirs(archive_dir, exist_ok=True)
    except OSError as exc:
        logger.error(
            "Failed to create weixin archive directory %s: %s",
            archive_dir,
            exc,
        )
        return
    for name in legacy_files:
        src = os.path.join(save_dir, name)
        new_name = _rewrite_weixin_in_session_filename(name)
        # ``new_name`` is non-None here (filtered above), but reassert for
        # the type checker.
        if new_name is None:
            continue
        dst = os.path.join(save_dir, new_name)
        archive_path = os.path.join(archive_dir, name)
        target_exists = os.path.exists(dst)
        try:
            if target_exists:
                # Canonical file already present: archive the legacy copy
                # and leave the live file untouched. ``shutil.move`` falls
                # back to copy+delete across filesystem boundaries.
                shutil.move(src, archive_path)
                logger.warning(
                    "Archived legacy weixin session file %s -> %s "
                    "(canonical %s already exists)",
                    src,
                    archive_path,
                    dst,
                )
            else:
                # Copy first, then archive the source. This keeps the
                # legacy file recoverable even if the move to ``dst`` is
                # interrupted.
                shutil.copy2(src, dst)
                shutil.move(src, archive_path)
                logger.warning(
                    "Migrated legacy weixin session file %s -> %s "
                    "(original archived to %s)",
                    src,
                    dst,
                    archive_path,
                )
        except OSError as exc:
            logger.error(
                "Failed to migrate session file %s -> %s: %s",
                src,
                dst,
                exc,
            )


def _rewrite_weixin_in_session_filename(name: str) -> str | None:
    """Return the canonical filename for a legacy weixin session file.

    File layout from ``_get_save_path`` is ``{safe_uid}_{safe_sid}.json``
    or ``{safe_sid}.json``. Returns ``None`` if the file does not match.

    NOTE: cannot use ``rsplit('_', 1)`` to find the boundary: WeChat
    user_ids contain ``_`` and session_ids end with ``@im.wechat``, so
    the rightmost ``_`` lives inside the session_id. Locate the literal
    ``_weixin--`` delimiter instead.
    """
    stem = name[: -len(".json")]
    delim = "_" + _LEGACY_WEIXIN_SAFE_PREFIX
    idx = stem.find(delim)
    if idx >= 0:
        safe_uid = stem[:idx]
        safe_sid_tail = stem[idx + len(delim) :]
        return (
            f"{safe_uid}_{_CANONICAL_WECHAT_SAFE_PREFIX}{safe_sid_tail}.json"
        )
    if stem.startswith(_LEGACY_WEIXIN_SAFE_PREFIX):
        return (
            _CANONICAL_WECHAT_SAFE_PREFIX
            + stem[len(_LEGACY_WEIXIN_SAFE_PREFIX) :]
            + ".json"
        )
    return None


class SafeJSONSession:
    """Filename-safe JSON session store with async file I/O."""

    def __init__(
        self,
        save_dir: str = "./",
    ) -> None:
        """Initialize the JSON session class.

        Args:
            save_dir (`str`, defaults to `"./"):
                The directory to save the session state.
        """
        self.save_dir = save_dir

    def _get_save_path(
        self,
        session_id: str,
        user_id: str,
        channel: str = "",
    ) -> str:
        """Return a filesystem-safe save path.

        Overrides the parent implementation to ensure the generated
        filename is valid on Windows, macOS and Linux.

        Args:
            session_id: Session identifier
            user_id: User identifier
            channel: Optional channel name for subdirectory separation

        Returns:
            Full path to the session file. If channel is provided,
            uses channels/{channel}/ subdirectory structure.
        """
        if not session_id:
            logger.error(
                "session_id is None or empty, cannot construct save path",
            )
            raise ValueError("session_id must not be None or empty")

        safe_sid = sanitize_filename(session_id)
        safe_uid = sanitize_filename(user_id) if user_id else ""

        # Guard against user_id == session_id (e.g. when upstream falls back
        # to session_id as user_id).  Duplicating the same token doubles the
        # filename length and can exceed Windows MAX_PATH (260 chars).
        if safe_uid and safe_uid == safe_sid:
            safe_uid = ""

        if safe_uid:
            filename = f"{safe_uid}_{safe_sid}.json"
        else:
            filename = f"{safe_sid}.json"

        if channel:
            safe_channel = sanitize_filename(channel)
            target_dir = os.path.join(self.save_dir, safe_channel)
            os.makedirs(target_dir, exist_ok=True)
            target_path = os.path.join(target_dir, filename)

            legacy_path = os.path.join(self.save_dir, filename)
            if not os.path.exists(target_path) and os.path.exists(legacy_path):
                try:
                    shutil.copy2(legacy_path, target_path)
                    logger.info(
                        "Migrated session file from %s to %s",
                        legacy_path,
                        target_path,
                    )
                except OSError as exc:
                    logger.warning(
                        "Failed to migrate session file %s to %s: %s",
                        legacy_path,
                        target_path,
                        exc,
                    )

            return target_path

        os.makedirs(self.save_dir, exist_ok=True)
        return os.path.join(self.save_dir, filename)

    async def save_session_state(
        self,
        session_id: str,
        user_id: str = "",
        channel: str = "",
        **state_modules_mapping,
    ) -> None:
        """Save state modules to a JSON file using async I/O."""
        state_dicts = {
            name: state_module.state_dict()
            for name, state_module in state_modules_mapping.items()
        }
        session_save_path = await run_sync_io(
            self._get_save_path,
            session_id,
            user_id,
            channel,
        )
        async with get_path_lock(session_save_path):
            await write_json_atomic_async(
                session_save_path,
                state_dicts,
                indent=None,
            )

        logger.info(
            "Saved session state to %s successfully.",
            session_save_path,
        )

    async def load_session_state(
        self,
        session_id: str,
        user_id: str = "",
        channel: str = "",
        allow_not_exist: bool = True,
        **state_modules_mapping,
    ) -> None:
        """Load state modules from a JSON file using async I/O."""
        session_save_path = await run_sync_io(
            self._get_save_path,
            session_id,
            user_id,
            channel,
        )
        try:
            states = await run_sync_io(
                _read_session_json,
                session_save_path,
            )
        except FileNotFoundError:
            states = None
        if states is not None:
            for name, state_module in state_modules_mapping.items():
                if name in states:
                    state_module.load_state_dict(states[name])
            logger.info(
                "Load session state from %s successfully.",
                session_save_path,
            )

        elif allow_not_exist:
            logger.info(
                "Session file %s does not exist. Skip loading session state.",
                session_save_path,
            )

        else:
            raise AgentStateError(
                session_id=session_id,
                message=(
                    f"Failed to load session state for file "
                    f"{session_save_path} because it does not exist"
                ),
            )

    async def update_session_state(
        self,
        session_id: str,
        key: Union[str, Sequence[str]],
        value,
        user_id: str = "",
        channel: str = "",
        create_if_not_exist: bool = True,
    ) -> None:
        session_save_path = await run_sync_io(
            self._get_save_path,
            session_id,
            user_id,
            channel,
        )

        path = key.split(".") if isinstance(key, str) else list(key)
        if not path:
            raise ConfigurationException(
                config_key="session.key",
                message="key path is empty",
            )

        async with get_path_lock(session_save_path):
            try:
                states = await run_sync_io(
                    _read_session_json,
                    session_save_path,
                )
            except FileNotFoundError as exc:
                if not create_if_not_exist:
                    raise AgentStateError(
                        session_id=session_id,
                        message=(
                            f"Session file {session_save_path}"
                            f" does not exist"
                        ),
                    ) from exc
                states = {}

            cur = states
            for k in path[:-1]:
                if k not in cur or not isinstance(cur[k], dict):
                    cur[k] = {}
                cur = cur[k]

            cur[path[-1]] = value

            await write_json_atomic_async(
                session_save_path,
                states,
                indent=None,
            )

        logger.info(
            "Updated session state key '%s' in %s successfully.",
            key,
            session_save_path,
        )

    async def get_session_state_dict(
        self,
        session_id: str,
        user_id: str = "",
        channel: str = "",
        allow_not_exist: bool = True,
    ) -> dict:
        """Return the session state dict from the JSON file.

        Args:
            session_id (`str`):
                The session id.
            user_id (`str`, default to `""`):
                The user ID for the storage.
            channel (`str`, default to `""`):
                The channel name for subdirectory separation.
            allow_not_exist (`bool`, defaults to `True`):
                Whether to allow the session to not exist. If `False`, raises
                an error if the session does not exist.

        Returns:
            `dict`:
                The session state dict loaded from the JSON file. Returns an
                empty dict if the file does not exist and
                `allow_not_exist=True`.
        """
        session_save_path = await run_sync_io(
            self._get_save_path,
            session_id,
            user_id,
            channel,
        )
        try:
            states = await run_sync_io(
                _read_session_json,
                session_save_path,
            )
        except FileNotFoundError:
            states = None
        if states is not None:
            logger.info(
                "Get session state dict from %s successfully.",
                session_save_path,
            )
            return states

        if allow_not_exist:
            logger.info(
                "Session file %s does not exist. Return empty state dict.",
                session_save_path,
            )
            return {}

        raise AgentStateError(
            session_id=session_id,
            message=(
                f"Failed to get session state for file {session_save_path} "
                f"because it does not exist"
            ),
        )
