from html import escape
import os
import uuid
from typing import Optional, Dict, List, Tuple, Any
from itsdangerous import URLSafeTimedSerializer, BadSignature
from micropie import App, HttpMiddleware, Request, SESSION_TIMEOUT


class SignedSessionMiddleware(HttpMiddleware):
    """Middleware to sign and verify session cookies using itsdangerous."""

    def __init__(self, app: App, secret_key: str, max_age: int = SESSION_TIMEOUT):
        self.app = app  # Store the App instance
        self.serializer = URLSafeTimedSerializer(secret_key)
        self.max_age = max_age

    async def before_request(self, request: Request) -> Optional[Dict]:
        """Verify the session_id cookie before processing the request."""
        cookies = self.app._parse_cookies(request.headers.get("cookie", ""))
        session_id = cookies.get("session_id", "")
        try:
            verified_id = self.serializer.loads(session_id, max_age=self.max_age)
            request.session = await self.app.session_backend.load(verified_id) or {}
        except BadSignature:
            request.session = {}  # Invalid or expired session_id
        return None

    async def after_request(
        self,
        request: Request,
        status_code: int,
        response_body: Any,
        extra_headers: List[Tuple[str, str]],
    ) -> Optional[Dict]:
        """Sign and set the session_id cookie after processing the request."""
        if request.session:
            cookies = self.app._parse_cookies(request.headers.get("cookie", ""))
            session_id = cookies.get("session_id", "") or str(uuid.uuid4())
            try:
                session_id = self.serializer.loads(session_id, max_age=self.max_age)
            except BadSignature:
                session_id = str(uuid.uuid4())
            signed_session_id = self.serializer.dumps(session_id)
            current_session = await self.app.session_backend.load(session_id) or {}
            if current_session != request.session:
                await self.app.session_backend.save(
                    session_id, request.session, SESSION_TIMEOUT
                )
            if not cookies.get("session_id"):
                extra_headers.append(
                    (
                        "Set-Cookie",
                        f"session_id={signed_session_id}; Path=/; SameSite=Lax; HttpOnly; Secure;",
                    )
                )
        return None


class Root(App):
    async def index(self):
        if "visits" not in self.request.session:
            self.request.session["visits"] = 1
        else:
            self.request.session["visits"] += 1
        visits = self.request.session["visits"]
        return f"You have visited {escape(str(visits))} times."


app = Root()
app.middlewares.append(SignedSessionMiddleware(app=app, secret_key="my-secret-key"))