"""Agent session management — memory between requests.

Maintains compressed conversation history per project, allowing the proxy
to avoid re-sending full context each time. Sessions are keyed by
(project_id, session_id) and stored in memory with optional persistence.
"""

from __future__ import annotations

import json
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path

import structlog

logger = structlog.get_logger()


@dataclass
class AgentSession:
    """Stateful session for an agent working on a project."""

    session_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
    project_id: str = ""
    metadata: dict = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)
    last_active: float = field(default_factory=time.time)

    # Compressed conversation history
    compressed_history: list[dict] = field(default_factory=list)

    # Active working set
    active_files: list[str] = field(default_factory=list)
    pending_tasks: list[str] = field(default_factory=list)

    # Token tracking
    total_tokens_sent: int = 0
    total_tokens_received: int = 0

    def touch(self) -> None:
        """Update last_active timestamp."""
        self.last_active = time.time()

    @property
    def id(self) -> str:
        return self.session_id

    @property
    def history(self) -> list[dict]:
        result = []
        user_entry = None
        for entry in self.compressed_history:
            if entry.get("role") == "user":
                user_entry = entry
                continue
            if entry.get("role") == "assistant" and user_entry is not None:
                result.append({
                    "user": user_entry.get("content", ""),
                    "assistant": entry.get("content", ""),
                    "metadata": entry.get("metadata", {}),
                    "ts": entry.get("ts", user_entry.get("ts", time.time())),
                })
                user_entry = None
        return result

    @property
    def message_count(self) -> int:
        return len(self.history)

    def add_exchange(
        self,
        user_msg: str | None = None,
        assistant_msg: str | None = None,
        metadata: dict | None = None,
        user_message: str | None = None,
        assistant_message: str | None = None,
    ) -> None:
        """Record a compressed user/assistant exchange."""
        user_text = user_message if user_message is not None else (user_msg or "")
        assistant_text = (
            assistant_message if assistant_message is not None else (assistant_msg or "")
        )
        self.compressed_history.append({
            "role": "user",
            "content": user_text[:500],
            "ts": time.time(),
        })
        self.compressed_history.append({
            "role": "assistant",
            "content": assistant_text[:500],
            "metadata": metadata or {},
            "ts": time.time(),
        })
        self.touch()

        # Keep history bounded
        if len(self.compressed_history) > 100:
            self.compressed_history = self.compressed_history[-60:]

    @property
    def age_seconds(self) -> float:
        return time.time() - self.created_at

    @property
    def idle_seconds(self) -> float:
        return time.time() - self.last_active


class SessionStore:
    """In-memory session store with TTL-based expiry.

    Parameters
    ----------
    max_idle_seconds : float
        Sessions inactive for longer than this are evicted.
    max_sessions : int
        Maximum number of concurrent sessions.
    """

    def __init__(
        self,
        base_dir: str | Path | None = None,
        ttl_seconds: float | None = None,
        max_idle_seconds: float = 3600,
        max_sessions: int = 100,
    ):
        self.base_dir = Path(base_dir) if base_dir is not None else None
        if self.base_dir is not None:
            self.base_dir.mkdir(parents=True, exist_ok=True)
        self.max_idle_seconds = ttl_seconds if ttl_seconds is not None else max_idle_seconds
        self.max_sessions = max_sessions
        self._sessions: dict[str, AgentSession] = {}
        self._load_sessions()

    def get_or_create(
        self,
        session_id: str = "",
        project_id: str = "",
    ) -> AgentSession:
        """Get an existing session or create a new one."""
        self._evict_expired()

        if session_id and session_id in self._sessions:
            session = self._sessions[session_id]
            session.touch()
            self._persist_session(session)
            return session

        session = AgentSession(
            session_id=session_id or uuid.uuid4().hex[:12],
            project_id=project_id,
        )
        self._sessions[session.session_id] = session
        self._persist_session(session)
        logger.info("session_created", session_id=session.session_id,
                     project_id=project_id)
        self._evict_expired()
        return session

    def get(self, session_id: str) -> AgentSession | None:
        """Get session by ID, or None if not found/expired."""
        session = self._sessions.get(session_id)
        if session and session.idle_seconds <= self.max_idle_seconds:
            return session
        return None

    def delete(self, session_id: str) -> bool:
        """Delete a session."""
        if session_id in self._sessions:
            del self._sessions[session_id]
            if self.base_dir is not None:
                path = self.base_dir / f"{session_id}.json"
                if path.exists():
                    path.unlink()
            return True
        return False

    def list_sessions(self) -> list[AgentSession]:
        """List all active sessions."""
        self._evict_expired()
        return list(self._sessions.values())

    def clear(self) -> None:
        session_ids = list(self._sessions.keys())
        for session_id in session_ids:
            self.delete(session_id)

    def _evict_expired(self) -> None:
        """Remove sessions that have been idle too long or exceed max count."""
        expired = [
            sid for sid, s in self._sessions.items()
            if s.idle_seconds > self.max_idle_seconds
        ]
        for sid in expired:
            self.delete(sid)

        # If still over limit, remove oldest
        if len(self._sessions) > self.max_sessions:
            by_age = sorted(self._sessions.items(), key=lambda x: x[1].last_active)
            for sid, _ in by_age[:len(self._sessions) - self.max_sessions]:
                self.delete(sid)

    def _session_path(self, session_id: str) -> Path | None:
        if self.base_dir is None:
            return None
        return self.base_dir / f"{session_id}.json"

    def _persist_session(self, session: AgentSession) -> None:
        path = self._session_path(session.session_id)
        if path is None:
            return
        path.write_text(json.dumps({
            "session_id": session.session_id,
            "project_id": session.project_id,
            "metadata": session.metadata,
            "created_at": session.created_at,
            "last_active": session.last_active,
            "compressed_history": session.compressed_history,
            "active_files": session.active_files,
            "pending_tasks": session.pending_tasks,
            "total_tokens_sent": session.total_tokens_sent,
            "total_tokens_received": session.total_tokens_received,
        }))

    def _load_sessions(self) -> None:
        if self.base_dir is None or not self.base_dir.exists():
            return
        for path in self.base_dir.glob("*.json"):
            try:
                data = json.loads(path.read_text())
                session = AgentSession(
                    session_id=data.get("session_id") or path.stem,
                    project_id=data.get("project_id", ""),
                    metadata=data.get("metadata", {}),
                    created_at=data.get("created_at", time.time()),
                    last_active=data.get("last_active", time.time()),
                    compressed_history=data.get("compressed_history", []),
                    active_files=data.get("active_files", []),
                    pending_tasks=data.get("pending_tasks", []),
                    total_tokens_sent=data.get("total_tokens_sent", 0),
                    total_tokens_received=data.get("total_tokens_received", 0),
                )
                self._sessions[session.session_id] = session
            except Exception:
                continue
