import base64
import binascii
import sqlite3
import cherrypy
from cherrypy import HTTPError, HTTPRedirect

from .config import Config
from .ssh_ca import SSHCA
from .utils import generate_access_token
from .database import Database, IntegrityError

from . import crypto

from oic.oic import Client as OICClient
from oic.utils.authn.client import CLIENT_AUTHN_METHOD
from oic.oic.message import AuthorizationResponse
from oic import rndstr

import functools
import logging
import os

from ipaddress import ip_network, ip_address
from typing import Any, Optional
from jinja2 import Environment, FileSystemLoader, select_autoescape

from natsort import natsorted

logger = logging.getLogger("server")

FILE_DIR = os.path.abspath(os.path.realpath(os.path.dirname(__file__)))
TEMPLATES_DIR = os.path.abspath(os.path.join(FILE_DIR, "..", "templates"))
STATIC_DIR = os.path.abspath(os.path.join(FILE_DIR, "..", "static"))

env = Environment(
    loader=FileSystemLoader(TEMPLATES_DIR), autoescape=select_autoescape()
)

SESSION_KEY_COOKIE = "SESSION_KEY"


class WAError(Exception):
    """An error that should be shown to the user."""

    def __init__(self, message, extra_info):
        super().__init__(message)
        self.message = message
        self.extra_info = extra_info


def get_token_from_header():
    """
    Gets the auth token from the HTTP Authorization header.
    If the header is missing, raises 401 Unauthorized.
    """
    if "Authorization" not in cherrypy.request.headers:
        raise HTTPError(401)

    hdr = cherrypy.request.headers["Authorization"]
    return hdr.removeprefix("Bearer ")


def get_token_from_cookie():
    """
    Gets the auth token from the SESSION_KEY cookie.
    If the cookie is missing, raises 401 Unauthorized.
    """
    if SESSION_KEY_COOKIE not in cherrypy.request.cookie:
        raise HTTPError(401)

    return cherrypy.request.cookie[SESSION_KEY_COOKIE].value


def render_on_error(template_name: str = "error.html"):
    """Render the given Jinja2 template on error."""

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            try:
                return fn(self, *args, **kwargs)
            except WAError as e:
                template = env.get_template(template_name)
                params = {"err": e.message, "extra_info": e.extra_info}
                params["wa_version"] = os.environ["WEBTECH_VERSION"]
                return template.render(params)

        return wrapper

    return decorator


def render_template(template_name: str):
    """Render the given Jinja2 template."""

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            template = env.get_template(template_name)
            params = fn(self, *args, **kwargs)
            params["wa_version"] = os.environ["WEBTECH_VERSION"]
            return template.render(params)

        return wrapper

    return decorator


def require_auth(as_cookie: bool = False):
    """
    Require authentication on this endpoint.

    This wrapper verifies that the request contains a valid authentication token.
    It then passes the authentication token to the wrapped function as the
    `auth_token` keyword parameter.

    If `as_cookie` is set to true, looks for an authentication token in the
    SESSION_KEY cookie.
    """

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(self, *args, **kwargs):
            auth_token: str
            if as_cookie:
                auth_token = get_token_from_cookie()
            else:
                auth_token = get_token_from_header()

            res = self.db.execute(
                "SELECT * FROM auth_tokens WHERE token = ?", (auth_token,)
            ).fetchone()

            if not res:
                if as_cookie:
                    # cookie has wrong token. erase cookie and redirect to /
                    cherrypy.response.cookie[SESSION_KEY_COOKIE] = auth_token
                    cherrypy.response.cookie[SESSION_KEY_COOKIE]["expires"] = 0
                    raise HTTPRedirect("/")
                else:
                    raise HTTPError(401)

            return fn(self, *args, **kwargs, auth_token=auth_token)

        return wrapper

    return decorator


def require_method(methods: Optional[list[str]] = None):
    """Require the use of a specific HTTP method. Default: GET only."""
    methods = methods or ["GET"]

    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            if cherrypy.request.method not in methods:
                raise HTTPError(405)
            return fn(*args, **kwargs)

        return wrapper

    return decorator


class Server:
    def __init__(self, config: Config, ssh_ca: SSHCA):
        self.config = config
        self.ssh_ca = ssh_ca

        db_path = self.config.lookup("general.db_path")
        self.db = Database(db_path)

        self._init_oic()

    def _init_oic(self):
        self.client_id = os.environ["OIDC_CLIENTID"]
        self.client_secret = os.environ["OIDC_SECRET"]
        self.oic_authority = os.environ["OIDC_AUTHORITY"]

        self.oic_client = OICClient(
            client_id=self.client_id, client_authn_method=CLIENT_AUTHN_METHOD
        )

        self.oic_client.provider_config(self.oic_authority)

    def _get_user_info(self, auth_token: str) -> Optional[dict[str, Any]]:
        """Return the user info for a given auth token from the db."""
        info = self.db.execute(
            "SELECT user_id, first_name, last_name, username FROM auth_tokens \
             JOIN users ON auth_tokens.user = users.user_id \
             WHERE auth_tokens.token = ?;",
            (auth_token,),
        ).fetchone()

        if not info:
            return None

        return dict(info)

    def _get_group_info(self, user_id: str) -> Optional[dict[str, str]]:
        info = self.db.execute(
            "SELECT name AS group_name FROM groups\
                         JOIN students ON students.group_name = groups.name\
                         JOIN users ON students.student_id = users.user_id\
                         WHERE users.user_id = ?;",
            (user_id,),
        ).fetchone()

        if not info:
            return None

        return dict(info)

    def _get_groups(self, user_id: str) -> Optional[list[str]]:
        student_info = self.db.execute(
            "SELECT group_name FROM students WHERE student_id = ?", (user_id,)
        ).fetchall()
        ta_info = self.db.execute(
            "SELECT group_name FROM tas WHERE ta = ?", (user_id,)
        ).fetchall()

        if student_info:
            return [row["group_name"] for row in student_info]

        if ta_info:
            return [row["group_name"] for row in ta_info]

    def _host_in_range(self, host) -> bool:
        """Checks whether a given host is in the ranges stored in the database."""
        addr = ip_address(host)
        ip_ranges = self.db.execute("SELECT ip_range FROM host_ranges").fetchall()
        for (ip_range,) in ip_ranges:
            net = ip_network(ip_range)
            if addr in net:
                return True

        return False

    def _has_active_session(self) -> bool:
        return SESSION_KEY_COOKIE in cherrypy.request.cookie

    def _do_login(self, state):
        params = {
            "client_id": self.client_id,
            "response_type": "code",
            "redirect_uri": "https://webtech-admin.datanose.nl/auth_callback",
            "scope": "openid",
            "state": f"{state}{rndstr()}",
        }

        auth_req = self.oic_client.construct_AuthorizationRequest(request_args=params)
        login_url = auth_req.request(self.oic_client.authorization_endpoint)

        raise HTTPRedirect(login_url)

    @cherrypy.expose
    def index(self):
        if self._has_active_session():
            return self.index_logged_in()  # type: ignore
        else:
            return self._do_login("frontend")

    @require_auth(as_cookie=True)
    @render_template("index.html")
    def index_logged_in(self, *, auth_token: str):
        info = self._get_user_info(auth_token)
        assert info

        group_info = self._get_group_info(info["user_id"]) or {}

        info.update(group_info)

        return info

    @cherrypy.expose
    def login_api(self):
        return self._do_login("api")

    @cherrypy.expose
    @render_on_error("error.html")
    def auth_callback(self, code, state):
        """
        Called on successful authentication.
        """
        # If you don't do this, do_access_token_request fails, because
        # this registers some state from the auth response.
        _ = self.oic_client.parse_response(
            AuthorizationResponse, info={"code": code, "state": state}, sformat="dict"
        )

        atr_args = {
            "code": code,
            "client_id": self.client_id,
            "client_secret": self.client_secret,
        }

        try:
            resp = self.oic_client.do_access_token_request(
                state=state,
                request_args=atr_args,
                authn_method="client_secret_post",
                scope="openid",
            )

            userinfo = self.oic_client.do_user_info_request(state=state)
            user_id = userinfo["uids"][0]
        except Exception as e:
            if state.startswith("frontend"):
                raise HTTPRedirect("/")
            else:
                raise HTTPError(400)

        logger.info(f"Got user ID: {user_id}")

        # Once we have the user info, generate an access token and store it.
        token = generate_access_token()
        try:
            self.db.execute(
                "INSERT INTO auth_tokens (user, token) VALUES (?, ?)", (user_id, token)
            )
        except IntegrityError:
            # Foreign key constraint failed, so user does not exist.
            if state.startswith("frontend"):
                raise WAError("unable to process login; user unknown", f"{user_id=}")
            else:
                raise HTTPError(422, f"Unknown user {user_id!r}")

        # Multiple options now; if this is a login for the API, redirect
        # to the webtech-ssh listener.

        # If this is a login for the front-end, set session cookie + redirect to
        # the index page.
        if state.startswith("frontend"):
            cherrypy.response.cookie[SESSION_KEY_COOKIE] = token
            cherrypy.response.cookie[SESSION_KEY_COOKIE]["secure"] = True
            raise HTTPRedirect("/")
        else:
            raise HTTPRedirect(f"http://localhost:19234/?token={token}")

    @cherrypy.expose
    @require_auth()
    @require_method(methods=["POST"])
    @cherrypy.tools.json_in()  # type: ignore
    def sign(self, *, auth_token: str):
        """Generate a certificate for a public key."""
        public_key = cherrypy.request.json.get("public_key")
        if public_key is None:
            raise HTTPError(400, "Missing public_key in POST body")

        user_info = self._get_user_info(auth_token)
        assert user_info

        try:
            cert = self.ssh_ca.do_sign(
                public_key.encode("ascii"),
                crypto.SSHCertificateType.USER,
                key_id=f"{user_info['first_name']} {user_info['last_name']} ({user_info['user_id']})",
                principals=(user_info["username"],),
            )

            return cert.public_bytes()
        except ValueError as e:
            raise HTTPError(400, "Unable to sign; invalid public key?") from e

    @cherrypy.expose
    def host_ca(self):
        """
        Return the host CA public key. This public key is trusted by
        clients, allowing them to connect to any host whose host keys
        are signed by us.
        """
        return self.ssh_ca._load_host_key().public.public_bytes(
            crypto.Encoding.OpenSSH, crypto.PublicFormat.OpenSSH
        )

    @cherrypy.expose
    def user_ca(self):
        """
        Return the user CA public key. This public key is trusted by
        servers, allowing users to log in when presenting a valid
        certificate signed by the user CA.
        """
        return self.ssh_ca._load_user_key().public.public_bytes(
            crypto.Encoding.OpenSSH, crypto.PublicFormat.OpenSSH
        )

    @cherrypy.expose
    @require_auth()
    @cherrypy.tools.json_out()  # type: ignore
    def user_info(self, *, auth_token: str):
        """Return the user info associated with the given auth token."""
        info = self._get_user_info(auth_token)
        assert info

        groups = self._get_groups(info["user_id"])

        if groups:
            info["groups"] = natsorted(groups)

        return dict(info)

    @require_method(methods=["POST"])
    @cherrypy.expose
    @cherrypy.tools.json_in()  # type: ignore
    def host_sign(self):
        """
        Sign a host key with the host CA key.
        Requires the host to be within the IP ranges specified in the db.
        """
        if "X-Forwarded-For" in cherrypy.request.headers:
            remote_ip = cherrypy.request.headers["X-Forwarded-For"]
        else:
            remote_ip = cherrypy.request.remote.ip

        if not self._host_in_range(remote_ip):
            raise HTTPError(401, "Not in valid range for host")

        host_key = cherrypy.request.json.get("host_key")
        host_names = cherrypy.request.json.get("host_names")
        if host_key is None or host_names is None:
            raise HTTPError(400, "Missing host_key or host_names in POST body")

        # We let the host pass the principal here because it is trusted anyway.
        try:
            cert = self.ssh_ca.do_sign(
                host_key.encode("ascii"),
                crypto.SSHCertificateType.HOST,
                key_id=f"{', '.join(host_names)} host certificate",
                principals=host_names,
            )

            return cert.public_bytes()
        except ValueError as e:
            raise HTTPError(400, "Unable to sign, invalid host key?") from e

    @require_auth()
    @require_method(methods=["GET"])
    @cherrypy.expose
    @cherrypy.tools.json_out()  # type: ignore
    def database_dump(self, *, auth_token: str):
        """
        Dump the current database as Base64.
        """
        info = self._get_user_info(auth_token)
        if not info:
            raise HTTPError(400, "Unable to get user info?")

        if info["user_id"].lower() != "skampen":
            raise HTTPError(401, "database dump not allowed")

        destination_db = sqlite3.connect(":memory:")
        self.db.conn.backup(destination_db)

        data = base64.b64encode(destination_db.serialize()).decode("ascii")
        return {"data": data}

    @require_auth()
    @require_method(methods=["POST"])
    @cherrypy.expose
    @cherrypy.tools.json_in()  # type: ignore
    @cherrypy.tools.json_out()  # type: ignore
    def database_slurp(self, *, auth_token: str):
        """
        Ingest a new database from base64-encoded data.
        """
        info = self._get_user_info(auth_token)
        if not info:
            raise HTTPError(400, "Unable to get user info?")

        if info["user_id"].lower() != "skampen":
            raise HTTPError(401, "database slurp not allowed")

        db_data = cherrypy.request.json.get("data")
        if db_data is None:
            raise HTTPError(400, "Missing data parameter")

        try:
            db_data = base64.b64decode(db_data)
        except binascii.Error:
            raise HTTPError(400, "Invalid base64 data")

        # validate that the data represents a sqlite db
        try:
            test_db = sqlite3.connect(":memory:")
            test_db.deserialize(db_data)
        except sqlite3.DatabaseError:
            raise HTTPError(400, "Invalid database!")

        self.db.conn.close()
        db_path = self.config.lookup("general.db_path")
        with open(db_path, "wb") as f:
            f.write(db_data)

        self.db = Database(db_path)

        return {"status": "ingested new database!"}

    def run(self):
        host = self.config.lookup("general.host")
        port = self.config.lookup("general.port")

        logger.info(f"Enabling static dir: {STATIC_DIR!r}")

        cherrypy.config.update(
            {
                "environment": "production",
                "server.socket_host": host,
                "server.socket_port": port,
            }
        )

        root_config = {
            "/static": {
                "tools.staticdir.on": True,
                "tools.staticdir.dir": STATIC_DIR,
            },
        }

        cherrypy.quickstart(self, "/", root_config)
