try:
    import jwt
    HAS_JWT = True
except ImportError:
    HAS_JWT = False
    try:
        import ubinascii
    except ImportError:
        import binascii as ubinascii
    try:
        import uhashlib as hashlib
    except ImportError:
        import hashlib
    try:
        import uhmac as hmac
    except ImportError:
        try:
            import hmac
        except ImportError:
            hmac = None
    import json

from microdot.microdot import invoke_handler
from microdot.helpers import wraps


class SessionDict(dict):
    """A session dictionary.

    The session dictionary is a standard Python dictionary that has been
    extended with convenience ``save()`` and ``delete()`` methods.
    """
    def __init__(self, request, session_dict):
        super().__init__(session_dict)
        self.request = request

    def save(self):
        """Update the session cookie."""
        self.request.app._session.update(self.request, self)

    def delete(self):
        """Delete the session cookie."""
        self.request.app._session.delete(self.request)


class Session:
    """Session handling

    :param app: The application instance.
    :param secret_key: The secret key, as a string or bytes object.
    :param cookie_options: A dictionary with cookie options to pass as
                           arguments to :meth:`Response.set_cookie()
                           <microdot.Response.set_cookie>`.
    """
    secret_key = None

    def __init__(self, app=None, secret_key=None, cookie_options=None):
        self.secret_key = secret_key
        self.cookie_options = cookie_options or {}
        if app is not None:
            self.initialize(app)

    def initialize(self, app, secret_key=None, cookie_options=None):
        if secret_key is not None:
            self.secret_key = secret_key
        if cookie_options is not None:
            self.cookie_options = cookie_options
        if 'path' not in self.cookie_options:
            self.cookie_options['path'] = '/'
        if 'http_only' not in self.cookie_options:
            self.cookie_options['http_only'] = True
        app._session = self

    def get(self, request):
        """Retrieve the user session.

        :param request: The client request.

        The return value is a session dictionary with the data stored in the
        user's session, or ``{}`` if the session data is not available or
        invalid.
        """
        if not self.secret_key:
            raise ValueError('The session secret key is not configured')
        if hasattr(request.g, '_session'):
            return request.g._session
        session = request.cookies.get('session')
        if session is None:
            request.g._session = SessionDict(request, {})
            return request.g._session
        request.g._session = SessionDict(request, self.decode(session))
        return request.g._session

    def update(self, request, session):
        """Update the user session.

        :param request: The client request.
        :param session: A dictionary with the update session data for the user.

        Applications would normally not call this method directly, instead they
        would use the :meth:`SessionDict.save` method on the session
        dictionary, which calls this method. For example::

            @app.route('/')
            @with_session
            def index(request, session):
                session['foo'] = 'bar'
                session.save()
                return 'Hello, World!'

        Calling this method adds a cookie with the updated session to the
        request currently being processed.
        """
        if not self.secret_key:
            raise ValueError('The session secret key is not configured')

        encoded_session = self.encode(session)

        @request.after_request
        def _update_session(request, response):
            response.set_cookie('session', encoded_session,
                                **self.cookie_options)
            return response

    def delete(self, request):
        """Remove the user session.

        :param request: The client request.

        Applications would normally not call this method directly, instead they
        would use the :meth:`SessionDict.delete` method on the session
        dictionary, which calls this method. For example::

            @app.route('/')
            @with_session
            def index(request, session):
                session.delete()
                return 'Hello, World!'

        Calling this method adds a cookie removal header to the request
        currently being processed.
        """
        @request.after_request
        def _delete_session(request, response):
            response.delete_cookie('session', **self.cookie_options)
            return response

    def encode(self, payload, secret_key=None):
        """Encode session data using JWT if available, otherwise use simple HMAC."""
        if HAS_JWT:
            return jwt.encode(payload, secret_key or self.secret_key,
                              algorithm='HS256')
        else:
            # Simple encoding for MicroPython: base64(json) + HMAC signature
            key = (secret_key or self.secret_key).encode() if isinstance(secret_key or self.secret_key, str) else (secret_key or self.secret_key)
            payload_json = json.dumps(payload)
            payload_b64 = ubinascii.b2a_base64(payload_json.encode()).decode().strip()
            
            # Create HMAC signature
            if hmac:
                # Use hmac module if available
                h = hmac.new(key, payload_json.encode(), hashlib.sha256)
            else:
                # Fallback: simple SHA256(key + message)
                h = hashlib.sha256(key + payload_json.encode())
            signature = ubinascii.b2a_base64(h.digest()).decode().strip()
            
            return f"{payload_b64}.{signature}"

    def decode(self, session, secret_key=None):
        """Decode session data using JWT if available, otherwise use simple HMAC."""
        if HAS_JWT:
            try:
                payload = jwt.decode(session, secret_key or self.secret_key,
                                     algorithms=['HS256'])
            except jwt.exceptions.PyJWTError:  # pragma: no cover
                return {}
            return payload
        else:
            try:
                # Simple decoding for MicroPython
                if '.' not in session:
                    return {}
                
                payload_b64, signature = session.rsplit('.', 1)
                payload_json = ubinascii.a2b_base64(payload_b64).decode()
                
                # Verify HMAC signature
                key = (secret_key or self.secret_key).encode() if isinstance(secret_key or self.secret_key, str) else (secret_key or self.secret_key)
                if hmac:
                    # Use hmac module if available
                    h = hmac.new(key, payload_json.encode(), hashlib.sha256)
                else:
                    # Fallback: simple SHA256(key + message)
                    h = hashlib.sha256(key + payload_json.encode())
                expected_signature = ubinascii.b2a_base64(h.digest()).decode().strip()
                
                if signature != expected_signature:
                    return {}
                
                return json.loads(payload_json)
            except Exception:
                return {}


def with_session(f):
    """Decorator that passes the user session to the route handler.

    The session dictionary is passed to the decorated function as an argument
    after the request object. Example::

        @app.route('/')
        @with_session
        def index(request, session):
            return 'Hello, World!'

    Note that the decorator does not save the session. To update the session,
    call the :func:`session.save() <microdot.session.SessionDict.save>` method.
    """
    @wraps(f)
    async def wrapper(request, *args, **kwargs):
        return await invoke_handler(
            f, request, request.app._session.get(request), *args, **kwargs)

    return wrapper
