import abc
import json
import logging
import logging.config
import sys
import time
from collections.abc import MutableMapping
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Dict, Final, Iterator, Optional, Union, cast, override

import trafaret as t
from aiohttp import web
from aiohttp.typedefs import Middleware

from ai.backend.logging import BraceStyleAdapter

log = BraceStyleAdapter(logging.getLogger("ai.backend.common.web.session"))

extra_config_headers = t.Dict({
    t.Key("X-BackendAI-Version", default=None): t.Null | t.String,
    t.Key("X-BackendAI-Encoded", default=None): t.Null | t.ToBool,
    t.Key("X-BackendAI-SessionID", default=None): t.Null | t.String,
}).allow_extra("*")

Handler = Callable[[web.Request], Awaitable[web.StreamResponse]]

if sys.version_info >= (3, 8):
    from typing import TypedDict
else:
    from typing_extensions import TypedDict


class _CookieParams(TypedDict, total=False):
    domain: Optional[str]
    max_age: Optional[int]
    path: str
    secure: Optional[bool]
    httponly: bool
    samesite: Optional[str]
    expires: Optional[str]


class SessionData(TypedDict, total=False):
    created: int
    expiration_dt: int
    session: Dict[str, Any]


DEFAULT_SESSION_AGE: Final[int] = 3600  # 1hour


def get_time() -> int:
    return int(time.time())


class Session(MutableMapping[str, Any]):
    """Session dict-like object."""

    def __init__(
        self,
        identity: Optional[Any],
        *,
        data: Optional[SessionData],
        new: bool,
        max_age: Optional[int] = None,
        lifespan: Optional[int] = None,
    ) -> None:
        self._changed: bool = False
        self._mapping: dict[str, Any] = {}
        self._identity = identity if data != {} else None
        self._new = new if data != {} else True
        self._max_age = max_age
        created = data.get("created", None) if data else None
        expiration_dt = data.get("expiration_dt") if data is not None else None
        session_data = data.get("session", None) if data else None
        now = get_time()
        is_expired = False
        if expiration_dt is not None:
            self._expiration_dt = expiration_dt
            if now > expiration_dt:
                is_expired = True
        else:
            age = now - created if created else now
            if max_age is not None:
                self._expiration_dt = now + max_age
                if age > max_age:
                    is_expired = True
            else:
                self._expiration_dt = now + DEFAULT_SESSION_AGE

        if self._new or created is None:
            self._created = now
        else:
            self._created = created

        if self._created > self._expiration_dt:
            self._expiration_dt = now + (lifespan or self.max_age or DEFAULT_SESSION_AGE)
            is_expired = False

        if is_expired:
            session_data = None

        if session_data is not None:
            self._mapping.update(session_data)

    def __repr__(self) -> str:
        return "<{} [new:{}, changed:{}, created:{} expiration:{}] {!r}>".format(
            self.__class__.__name__,
            self.new,
            self._changed,
            self.created,
            self.expiration_dt,
            self._mapping,
        )

    @property
    def new(self) -> bool:
        return self._new

    @property
    def identity(self) -> Optional[Any]:  # type: ignore[misc]
        return self._identity

    @property
    def created(self) -> int:
        return self._created

    def update_creation_time(self) -> None:
        self._created = get_time()
        self._changed = True

    @property
    def empty(self) -> bool:
        return not bool(self._mapping)

    @property
    def expiration_dt(self) -> int:
        return self._expiration_dt

    @expiration_dt.setter
    def expiration_dt(self, value: int) -> None:
        self._expiration_dt = value
        self._changed = True

    @property
    def max_age(self) -> Optional[int]:
        return self._max_age

    @max_age.setter
    def max_age(self, value: Optional[int]) -> None:
        self._max_age = value

    def changed(self) -> None:
        self._changed = True

    def invalidate(self) -> None:
        self._changed = True
        self._mapping = {}

    def set_new_identity(self, identity: Optional[Any]) -> None:
        if not self._new:
            raise RuntimeError("Can't change identity for a session which is not new")

        self._identity = identity

    def __len__(self) -> int:
        return len(self._mapping)

    def __iter__(self) -> Iterator[str]:
        return iter(self._mapping)

    def __contains__(self, key: object) -> bool:
        return key in self._mapping

    def __getitem__(self, key: str) -> Any:
        return self._mapping[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self._mapping[key] = value
        self._changed = True
        self._created = get_time()

    def __delitem__(self, key: str) -> None:
        del self._mapping[key]
        self._changed = True
        self._created = get_time()


SESSION_KEY = "aiohttp_session"
STORAGE_KEY = "aiohttp_session_storage"


async def get_session(request: web.Request) -> Session:
    session = request.get(SESSION_KEY)
    if session is None:
        storage = request.get(STORAGE_KEY)
        if storage is None:
            raise RuntimeError("Install aiohttp_session middleware in your aiohttp.web.Application")

        session = await storage.load_session(request)
        if not isinstance(session, Session):
            raise RuntimeError(
                "Installed {!r} storage should return session instance "
                "on .load_session() call, got {!r}.".format(storage, session)
            )
        request[SESSION_KEY] = session
    return session


async def new_session(request: web.Request) -> Session:
    storage = request.get(STORAGE_KEY)
    if storage is None:
        raise RuntimeError("Install aiohttp_session middleware in your aiohttp.web.Application")

    session = await storage.new_session()
    if not isinstance(session, Session):
        raise RuntimeError(
            "Installed {!r} storage should return session instance "
            "on .load_session() call, got {!r}.".format(storage, session)
        )
    request[SESSION_KEY] = session
    return session


def session_middleware(storage: "AbstractStorage") -> Middleware:
    if not isinstance(storage, AbstractStorage):
        raise RuntimeError(f"Expected AbstractStorage got {storage}")

    @web.middleware
    async def factory(request: web.Request, handler: Handler) -> web.StreamResponse:
        request[STORAGE_KEY] = storage
        raise_response = False
        # TODO aiohttp 4:
        # Remove Union from response, and drop the raise_response variable
        response: Union[web.StreamResponse, web.HTTPException]
        try:
            response = await handler(request)
        except web.HTTPException as exc:
            response = exc
            raise_response = True
        if not isinstance(response, (web.StreamResponse, web.HTTPException)):
            raise RuntimeError(f"Expect response, not {type(response)!r}")
        if not isinstance(response, (web.Response, web.HTTPException)):
            # likely got websocket or streaming
            return response
        if response.prepared:
            raise RuntimeError("Cannot save session data into prepared response")
        session = request.get(SESSION_KEY)
        if session is not None:
            if session._changed:
                await storage.save_session(request=request, response=response, session=session)
        if raise_response:
            raise cast(web.HTTPException, response)
        return response

    return factory


def setup(app: web.Application, storage: "AbstractStorage") -> None:
    """Setup the library in aiohttp fashion."""

    app.middlewares.append(session_middleware(storage))


class AbstractStorage(metaclass=abc.ABCMeta):
    def __init__(
        self,
        *,
        cookie_name: str = "AIOHTTP_SESSION",
        domain: Optional[str] = None,
        max_age: Optional[int] = None,
        path: str = "/",
        secure: Optional[bool] = None,
        httponly: bool = True,
        samesite: Optional[str] = None,
        encoder: Callable[[object], str] = json.dumps,
        decoder: Callable[[str], Any] = json.loads,
    ) -> None:
        self._cookie_name = cookie_name
        self._cookie_params = _CookieParams(
            domain=domain,
            max_age=max_age,
            path=path,
            secure=secure,
            httponly=httponly,
            samesite=samesite,
        )
        self._max_age = max_age
        self._encoder = encoder
        self._decoder = decoder

    @property
    def cookie_name(self) -> str:
        return self._cookie_name

    @property
    def max_age(self) -> Optional[int]:
        return self._max_age

    @property
    def cookie_params(self) -> _CookieParams:
        return self._cookie_params

    def _get_session_data(self, session: Session) -> SessionData:
        if session.empty:
            return {}
        return {
            "created": session.created,
            "session": session._mapping,
            "expiration_dt": session.expiration_dt,
        }

    async def new_session(self) -> Session:
        return Session(
            None,
            data=None,
            new=True,
            max_age=self.max_age,
        )

    @abc.abstractmethod
    async def load_session(self, request: web.Request) -> Session:
        pass

    @abc.abstractmethod
    async def save_session(
        self,
        request: web.Request,
        response: web.StreamResponse,
        session: Session,
        session_extension: int | None = None,
    ) -> None:
        pass

    def load_cookie(self, request: web.Request) -> Optional[str]:
        # TODO: Remove explicit type annotation when aiohttp 3.8 is out
        cookie: Optional[str] = request.cookies.get(self._cookie_name)
        return cookie

    def save_cookie(
        self,
        response: web.StreamResponse,
        cookie_data: str,
        *,
        expiration_dt: float | None = None,
        max_age: int | None = None,
    ) -> None:
        params = self._cookie_params.copy()
        if max_age is not None:
            params["max_age"] = max_age
        if expiration_dt is not None:
            params["expires"] = datetime.fromtimestamp(expiration_dt, tz=timezone.utc).isoformat()
        if not cookie_data:
            # Ignoring type for params until aiohttp#4238 is released
            response.del_cookie(
                name=self._cookie_name, domain=params["domain"], path=params["path"]
            )
        else:
            response.set_cookie(name=self._cookie_name, value=cookie_data, **params)


class SimpleCookieStorage(AbstractStorage):
    """Simple JSON storage.

    Doesn't any encryption/validation, use it for tests only"""

    def __init__(
        self,
        *,
        cookie_name: str = "AIOHTTP_SESSION",
        domain: Optional[str] = None,
        max_age: Optional[int] = None,
        path: str = "/",
        secure: Optional[bool] = None,
        httponly: bool = True,
        samesite: Optional[str] = None,
        encoder: Callable[[object], str] = json.dumps,
        decoder: Callable[[str], Any] = json.loads,
    ) -> None:
        super().__init__(
            cookie_name=cookie_name,
            domain=domain,
            max_age=max_age,
            path=path,
            secure=secure,
            httponly=httponly,
            samesite=samesite,
            encoder=encoder,
            decoder=decoder,
        )

    async def load_session(self, request: web.Request) -> Session:
        cookie = self.load_cookie(request)
        if cookie is None:
            return Session(None, data=None, new=True, max_age=self.max_age)

        data = self._decoder(cookie)
        return Session(None, data=data, new=False, max_age=self.max_age)

    @override
    async def save_session(
        self,
        request: web.Request,
        response: web.StreamResponse,
        session: Session,
        session_extension: int | None = None,
    ) -> None:
        cookie_data = self._encoder(self._get_session_data(session))
        if session_extension is not None:
            session.expiration_dt = get_time() + session_extension
        self.save_cookie(
            response, cookie_data, max_age=session.max_age, expiration_dt=session.expiration_dt
        )
