from __future__ import annotations

from contextlib import contextmanager
import json
from pathlib import Path
from typing import Callable, Iterator

from filelock import FileLock

from autosana_cli.models import Session, STATE_VERSION, autosana_home


class StateStore:
    def __init__(
        self,
        *,
        root_dir: Path | None = None,
        state_file: Path | None = None,
        lock_file: Path | None = None,
    ) -> None:
        self.root_dir = root_dir or autosana_home()
        self.state_file = state_file or (self.root_dir / "state.json")
        self.lock_file = lock_file or (self.root_dir / "state.lock")

    def ensure_layout(self) -> None:
        self.root_dir.mkdir(parents=True, exist_ok=True)
        (self.root_dir / "logs").mkdir(parents=True, exist_ok=True)
        self.lock_file.parent.mkdir(parents=True, exist_ok=True)

    @contextmanager
    def locked(self) -> Iterator[None]:
        self.ensure_layout()
        with FileLock(str(self.lock_file)):
            yield

    def load_sessions(self) -> list[Session]:
        if not self.state_file.exists():
            return []
        with self.state_file.open("r", encoding="utf-8") as handle:
            raw = json.load(handle)
        return [Session.from_dict(item) for item in raw.get("sessions", [])]

    def save_sessions(self, sessions: list[Session]) -> None:
        self.ensure_layout()
        payload = {"version": STATE_VERSION, "sessions": [session.to_dict() for session in sessions]}
        with self.state_file.open("w", encoding="utf-8") as handle:
            json.dump(payload, handle, indent=2)
            handle.write("\n")

    def read_locked(self) -> list[Session]:
        with self.locked():
            return self.load_sessions()

    def write_locked(self, sessions: list[Session]) -> None:
        with self.locked():
            self.save_sessions(sessions)

    def update_locked(self, updater: Callable[[list[Session]], list[Session]]) -> list[Session]:
        with self.locked():
            sessions = self.load_sessions()
            updated = updater(sessions)
            self.save_sessions(updated)
            return updated

    def add_session(self, session: Session) -> list[Session]:
        def updater(sessions: list[Session]) -> list[Session]:
            return [existing for existing in sessions if existing.session_id != session.session_id] + [session]

        return self.update_locked(updater)

    def remove_session(self, session_id: str) -> tuple[list[Session], list[Session]]:
        removed: list[Session] = []

        def updater(sessions: list[Session]) -> list[Session]:
            nonlocal removed
            keep = [session for session in sessions if session.session_id != session_id]
            removed = [session for session in sessions if session.session_id == session_id]
            return keep

        return self.update_locked(updater), removed

    def remove_matching(self, predicate: Callable[[Session], bool]) -> tuple[list[Session], list[Session]]:
        removed: list[Session] = []

        def updater(sessions: list[Session]) -> list[Session]:
            nonlocal removed
            removed = [session for session in sessions if predicate(session)]
            return [session for session in sessions if not predicate(session)]

        return self.update_locked(updater), removed

    def prune_dead_sessions(self, is_session_alive: Callable[[Session], bool]) -> tuple[list[Session], list[Session]]:
        removed: list[Session] = []

        def updater(sessions: list[Session]) -> list[Session]:
            nonlocal removed
            alive: list[Session] = []
            for session in sessions:
                if is_session_alive(session):
                    alive.append(session)
                else:
                    removed.append(session)
            return alive

        return self.update_locked(updater), removed
