import json
import re
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

_SAFE_ID_RE = re.compile(r"[^a-zA-Z0-9_.\-]")

_SESSION_FORMAT_VERSION = 2
_TOOL_OUTPUT_STUB = "[tool output processed]"
_SKILL_TOOL_NAME = "use_skill"
_SKILL_CONTENT_PREFIX = "__skill_content__:"


@dataclass
class Session:
    chat_id: str
    messages: list[dict]
    created_at: str
    updated_at: str
    bot_id: str | None = None
    records: list[dict] = field(default_factory=list)


class SessionStore:
    def __init__(self, sessions_dir: Path):
        self._sessions_dir = sessions_dir
        self._sessions_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def load(self, chat_id: str, **kwargs) -> Session:
        """Alias for load_full — returns full session with all messages.

        Accepts and ignores bot_id for backward compatibility.
        """
        return self.load_full(chat_id)

    def load_full(self, chat_id: str, **kwargs) -> Session:
        """Load a session from JSONL, returning an empty one when missing.

        Returns ALL messages from the on-disk file (full history).
        Accepts and ignores bot_id for backward compatibility.
        """
        path = self._session_path(chat_id)
        if not path.exists():
            now = _now_iso()
            return Session(
                chat_id=chat_id,
                messages=[],
                created_at=now,
                updated_at=now,
            )

        lines = _read_non_empty_lines(path)
        if not lines:
            now = _now_iso()
            return Session(
                chat_id=chat_id,
                messages=[],
                created_at=now,
                updated_at=now,
            )

        first = _parse_line(path, lines[0], line_number=1)
        if not isinstance(first, dict) or first.get("type") != "meta":
            raise ValueError(f"Invalid session format in {path}: expected v2 meta record on line 1")
        version = first.get("version")
        try:
            parsed_version = int(version)
        except (TypeError, ValueError) as exc:
            raise ValueError(f"Invalid session version in {path}: {version!r}") from exc
        if parsed_version != _SESSION_FORMAT_VERSION:
            raise ValueError(
                f"Unsupported session version in {path}: expected"
                f" {_SESSION_FORMAT_VERSION}, got {version}"
            )
        return _load_v2(path, lines, meta=first, chat_id=chat_id)

    def load_for_llm(
        self,
        chat_id: str,
        *,
        tool_results_keep: int = 4,
        **kwargs,
    ) -> list[dict]:
        """Build LLM context from on-disk session.

        Uses the latest summary as a base, appends messages after it,
        prunes old tool results, and compacts skill content.
        Full on-disk history is never modified.
        Accepts and ignores bot_id for backward compatibility.
        """
        session = self.load_full(chat_id)
        messages = list(session.messages)

        # Compact skill tool results
        messages = _compact_skill_results(messages)

        # Prune old tool results
        messages = _prune_tool_results(messages, keep_turns=tool_results_keep)

        return messages

    def save(self, session: Session, **kwargs) -> None:
        """Persist a session to JSONL v2 format.

        Accepts and ignores bot_id for backward compatibility.
        """
        path = self._session_path(session.chat_id)
        previous_updated_at = session.updated_at
        session.updated_at = _now_iso()
        if not session.created_at or (
            not path.exists() and session.created_at == previous_updated_at
        ):
            session.created_at = session.updated_at
        created_at = session.created_at
        updated_at = session.updated_at

        meta: dict[str, Any] = {
            "type": "meta",
            "version": _SESSION_FORMAT_VERSION,
            "chat_id": session.chat_id,
            "created_at": created_at,
            "updated_at": updated_at,
        }

        with path.open("w", encoding="utf-8") as file:
            file.write(json.dumps(meta, ensure_ascii=False) + "\n")
            for idx, message in enumerate(session.messages, start=1):
                record = _msg_record(message, msg_id=message.get("_msg_id", idx))
                file.write(json.dumps(record, ensure_ascii=False) + "\n")
            # Write any summary records
            for rec in session.records:
                if rec.get("type") == "summary":
                    file.write(json.dumps(rec, ensure_ascii=False) + "\n")

    def append(self, chat_id: str, message: dict, **kwargs) -> None:
        """Append one message, creating a new session file when needed.

        Accepts and ignores bot_id for backward compatibility.
        """
        path = self._session_path(chat_id)
        now = _now_iso()

        if not path.exists():
            self.save(
                Session(
                    chat_id=chat_id,
                    messages=[message],
                    created_at=now,
                    updated_at=now,
                ),
            )
            return

        session = self.load(chat_id)
        session.chat_id = chat_id
        if not session.created_at:
            session.created_at = now
        session.updated_at = now
        session.messages.append(message)
        self.save(session)

    def list_sessions(self, **kwargs) -> list[str]:
        """Return all chat IDs that have persisted sessions.

        Accepts and ignores bot_id for backward compatibility.
        """
        chat_ids: list[str] = []
        for path in self._sessions_dir.glob("*.jsonl"):
            if not path.is_file():
                continue
            stem = path.name.removesuffix(".jsonl")
            chat_ids.append(stem)
        return sorted(chat_ids)

    def delete(self, chat_id: str, **kwargs) -> None:
        """Delete a persisted session file.

        Accepts and ignores bot_id for backward compatibility.
        """
        self._session_path(chat_id).unlink(missing_ok=True)

    def rotate(self, chat_id: str) -> str | None:
        """Archive current session file with a compact UTC timestamp."""
        path = self._session_path(chat_id)
        if not path.exists() or path.stat().st_size == 0:
            return None

        ts = datetime.now(UTC).strftime("%Y%m%dT%H%M%S")
        safe_id = self._safe_filename(chat_id)
        archive_path = self._sessions_dir / f"{safe_id}.{ts}.jsonl"
        path.rename(archive_path)
        return ts

    def list_archived(self, chat_id: str) -> list[dict]:
        """List archived sessions for a chat, newest first (max 20)."""
        entries: list[dict] = []
        safe_id = self._safe_filename(chat_id)
        prefix = f"{safe_id}."

        for path in self._sessions_dir.glob(f"{safe_id}.*.jsonl"):
            if not path.is_file() or path.name == f"{safe_id}.jsonl":
                continue
            if not path.name.startswith(prefix):
                continue

            ts = path.name[len(prefix) : -len(".jsonl")]
            if not ts:
                continue

            lines = _read_non_empty_lines(path)
            msg_count = max(len(lines) - 1, 0)
            preview = ""
            for raw_line in lines[1:]:
                try:
                    record = json.loads(raw_line)
                except json.JSONDecodeError:  # pragma: no cover - defensive path
                    continue
                if (
                    isinstance(record, dict)
                    and record.get("type") == "msg"
                    and record.get("role") == "user"
                ):
                    preview = str(record.get("content", ""))[:60]
                    break

            entries.append(
                {
                    "ts": ts,
                    "msg_count": msg_count,
                    "preview": preview,
                    "path": path,
                }
            )

        entries.sort(key=lambda item: item["ts"], reverse=True)
        return entries[:20]

    def restore(self, chat_id: str, archive_ts: str) -> bool:
        """Restore an archived session into the active session path."""
        safe_id = self._safe_filename(chat_id)
        safe_ts = self._safe_filename(archive_ts)
        archive_path = self._sessions_dir / f"{safe_id}.{safe_ts}.jsonl"
        if not archive_path.exists():
            return False

        current_path = self._session_path(chat_id)
        if current_path.exists() and current_path.stat().st_size > 0:
            self.rotate(chat_id)

        archive_path.rename(current_path)
        return True

    @staticmethod
    def _safe_filename(value: str) -> str:
        """Sanitize a value for safe use in filenames (no path traversal)."""
        return _SAFE_ID_RE.sub("_", value)

    def _session_path(self, chat_id: str) -> Path:
        safe_id = self._safe_filename(chat_id)
        path = (self._sessions_dir / f"{safe_id}.jsonl").resolve()
        if not path.is_relative_to(self._sessions_dir.resolve()):
            raise ValueError(f"Invalid chat_id: {chat_id!r}")
        return path


# ------------------------------------------------------------------
# v2 record helpers
# ------------------------------------------------------------------


def _msg_record(message: dict, *, msg_id: int) -> dict:
    """Build a v2 msg record from a message dict."""
    content = message.get("content")
    if (
        message.get("role") == "tool"
        and message.get("name") == _SKILL_TOOL_NAME
        and "content" in message
    ):
        content = _persisted_skill_placeholder(content)

    record: dict[str, Any] = {
        "type": "msg",
        "id": msg_id,
        "role": message.get("role", "user"),
        "ts": message.get("ts", _now_iso()),
    }
    if "content" in message:
        record["content"] = content
    if "tool_calls" in message:
        record["tool_calls"] = message["tool_calls"]
    if "tool_call_id" in message:
        record["tool_call_id"] = message["tool_call_id"]
    if "name" in message:
        record["name"] = message["name"]
    return record


def _record_to_message(record: dict) -> dict:
    """Convert a v2 msg record back to a plain message dict."""
    msg: dict[str, Any] = {"role": record.get("role", "user")}
    if "content" in record:
        msg["content"] = record["content"]
    if "tool_calls" in record:
        msg["tool_calls"] = record["tool_calls"]
    if "tool_call_id" in record:
        msg["tool_call_id"] = record["tool_call_id"]
    if "name" in record:
        msg["name"] = record["name"]
    # Stash msg_id for load_for_llm filtering
    if "id" in record:
        msg["_msg_id"] = record["id"]
    return msg


# ------------------------------------------------------------------
# v2 loader
# ------------------------------------------------------------------


def _load_v2(
    path: Path,
    lines: list[str],
    *,
    meta: dict,
    chat_id: str,
    bot_id: str | None = None,
) -> Session:
    created_at = str(meta.get("created_at") or _now_iso())
    updated_at = str(meta.get("updated_at") or created_at)
    resolved_chat_id = str(meta.get("chat_id") or chat_id)

    messages: list[dict] = []
    records: list[dict] = []

    for index, raw_line in enumerate(lines[1:], start=2):
        record = _parse_line(path, raw_line, line_number=index)
        if not isinstance(record, dict):
            raise ValueError(f"Invalid session record in {path} on line {index}: expected object")
        rtype = record.get("type")
        if rtype == "msg":
            messages.append(_record_to_message(record))
        elif rtype == "summary":
            records.append(record)
        else:
            raise ValueError(f"Invalid session record type in {path} on line {index}: {rtype!r}")

    return Session(
        chat_id=resolved_chat_id,
        messages=messages,
        created_at=created_at,
        updated_at=updated_at,
        records=records,
    )


# ------------------------------------------------------------------
# LLM context building helpers
# ------------------------------------------------------------------


def _persisted_skill_placeholder(content: Any) -> str:
    existing_name = _extract_skill_placeholder_name(content)
    if existing_name is not None:
        return f"[skill:{existing_name} loaded]"

    text = str(content)
    if text.startswith(_SKILL_CONTENT_PREFIX):
        first_line = text.split("\n", 1)[0]
        skill_name = first_line[len(_SKILL_CONTENT_PREFIX) :].strip() or "unknown"
        return f"[skill:{skill_name} loaded]"
    return "[skill:unknown loaded]"


def _extract_skill_placeholder_name(content: Any) -> str | None:
    text = str(content).strip()
    prefix = "[skill:"
    suffix = " loaded]"
    if text.startswith(prefix) and text.endswith(suffix):
        name = text[len(prefix) : -len(suffix)].strip()
        if name:
            return name
    return None


def _prune_tool_results(messages: list[dict], *, keep_turns: int = 4) -> list[dict]:
    """Replace tool result content with a stub for turns older than keep_turns.

    A 'turn' is counted as each user message from the end.
    """
    if keep_turns < 0:
        return [_strip_internal(m) for m in messages]

    # Find user message indices to identify turn boundaries
    user_indices: list[int] = [i for i, m in enumerate(messages) if m.get("role") == "user"]

    if len(user_indices) <= keep_turns:
        return [_strip_internal(m) for m in messages]

    # Messages before the Nth-from-last user message are eligible for pruning
    cutoff_index = user_indices[-keep_turns]

    result: list[dict] = []
    for i, msg in enumerate(messages):
        clean = _strip_internal(msg)
        if i < cutoff_index and msg.get("role") == "tool":
            clean["content"] = _TOOL_OUTPUT_STUB
        result.append(clean)
    return result


def _strip_internal(msg: dict) -> dict:
    """Remove internal metadata keys before sending to LLM."""
    clean = dict(msg)
    clean.pop("_msg_id", None)
    return clean


def _compact_skill_results(messages: list[dict]) -> list[dict]:
    """Replace use_skill tool result content with a compact stub."""
    # Build a map of tool_call_id -> skill name from assistant tool_calls
    skill_call_names: dict[str, str] = {}
    for msg in messages:
        if msg.get("role") == "assistant" and "tool_calls" in msg:
            for tc in msg["tool_calls"]:
                fn = tc.get("function", {})
                if fn.get("name") == _SKILL_TOOL_NAME:
                    call_id = tc.get("id", "")
                    try:
                        args = json.loads(fn.get("arguments", "{}"))
                        skill_name = args.get("name", "unknown")
                    except (json.JSONDecodeError, TypeError):
                        skill_name = "unknown"
                    skill_call_names[call_id] = skill_name

    result: list[dict] = []
    for msg in messages:
        if msg.get("role") == "tool":
            content = str(msg.get("content", ""))
            # Detect via tool_call chain
            if (
                skill_call_names
                and msg.get("name") == _SKILL_TOOL_NAME
                and msg.get("tool_call_id") in skill_call_names
            ):
                compacted = dict(msg)
                skill_name = skill_call_names[msg["tool_call_id"]]
                compacted["content"] = f"[skill:{skill_name} loaded]"
                result.append(compacted)
            # Detect via prefix tag (fallback for collapsed history)
            elif content.startswith(_SKILL_CONTENT_PREFIX):
                compacted = dict(msg)
                # Extract skill name from prefix line
                first_line = content.split("\n", 1)[0]
                skill_name = first_line[len(_SKILL_CONTENT_PREFIX) :]
                compacted["content"] = f"[skill:{skill_name} loaded]"
                result.append(compacted)
            else:
                result.append(msg)
        else:
            result.append(msg)
    return result


# ------------------------------------------------------------------
# Shared utilities
# ------------------------------------------------------------------


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


def _read_non_empty_lines(path: Path) -> list[str]:
    with path.open("r", encoding="utf-8") as file:
        return [line.strip() for line in file if line.strip()]


def _parse_line(path: Path, raw_line: str, *, line_number: int):
    try:
        return json.loads(raw_line)
    except json.JSONDecodeError as exc:  # pragma: no cover - defensive path
        raise ValueError(f"Invalid JSON in {path} on line {line_number}") from exc
