from .models import Server
from contextlib import AbstractContextManager
from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import func
from typing import Callable, Iterator


class ServerRepository:
    def __init__(
        self, session_factory: Callable[..., AbstractContextManager[Session]]
    ) -> None:
        self.session_factory = session_factory

    def get_all(self) -> Iterator[Server]:
        with self.session_factory() as session:
            return session.query(Server).all()

    def get_by_address(self, server_address: str) -> Server:
        with self.session_factory() as session:
            server = (
                session.query(Server).filter(Server.address == server_address).first()
            )
            if not server:
                raise ServerNotFoundError(server_address)
            return server

    def add(self, server_address: str, json: str) -> Server:
        with self.session_factory() as session:
            server = Server(address=server_address, json=json)
            session.add(server)
            session.commit()
            session.refresh(server)
            return server

    def delete_by_address(self, server_address: int) -> None:
        with self.session_factory() as session:
            entity: Server = (
                session.query(Server).filter(Server.address == server_address).first()
            )
            if not entity:
                raise ServerNotFoundError(server_address)
            session.delete(entity)
            session.commit()

    def get_random(self) -> Server:
        with self.session_factory() as session:
            return session.query(Server).order_by(func.rand()).limit(1).first()


class NotFoundError(Exception):
    entity_name: str

    def __init__(self, entity_id):
        super().__init__(f"{self.entity_name} not found, id: {entity_id}")


class ServerNotFoundError(NotFoundError):
    entity_name: str = "Server"
