#
#   Copyright © 2021,2022 IsardVDI S.L.
#   Copyright © 2022 Evilham <contact@evilham.com>
#
#   This file is part of DD
#
#   DD is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or (at your
#   option) any later version.
#
#   DD is distributed in the hope that it will be useful, but WITHOUT ANY
#   WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
#   FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
#   details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with DD. If not, see <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: AGPL-3.0-or-later

import json
import logging as log
import os
import time
import traceback
from datetime import datetime, timedelta
from pprint import pprint

import yaml
from jinja2 import Environment, FileSystemLoader
from keycloak import KeycloakAdmin

from .api_exceptions import Error
from .helpers import get_recursive_groups, kpath2kpaths
from .postgres import Postgres

from typing import cast, Any, Dict, Iterable, List, Optional

DDUser = Dict[str, Any]

# TODO: Improve typing of these class and simplify it

class KeycloakClient:
    """https://www.keycloak.org/docs-api/13.0/rest-api/index.html
    https://github.com/marcospereirampj/python-keycloak
    https://gist.github.com/kaqfa/99829941121188d7cef8271f93f52f1f
    """
    url : str
    username : str
    password : str
    realm : str
    verify : bool
    keycloak_pg : Postgres
    keycloak_admin : KeycloakAdmin

    def __init__(
        self,
        url : str="http://dd-sso-keycloak:8080/auth/",
        username : str=os.environ.get("KEYCLOAK_USER", ""),
        password : str=os.environ.get("KEYCLOAK_PASSWORD", ""),
        realm : str="master",
        verify : bool=True,
    ) -> None:
        self.url = url
        self.username = username
        self.password = password
        self.realm = realm
        self.verify = verify

        self.keycloak_pg = Postgres(
            "dd-apps-postgresql",
            "keycloak",
            os.environ.get("KEYCLOAK_DB_USER", ""),
            os.environ.get("KEYCLOAK_DB_PASSWORD", ""),
        )

    def connect(self) -> None:
        self.keycloak_admin = KeycloakAdmin(
            server_url=self.url,
            username=self.username,
            password=self.password,
            realm_name=self.realm,
            verify=self.verify,
        )

    # from keycloak import KeycloakAdmin
    # keycloak_admin = KeycloakAdmin(server_url="http://dd-sso-keycloak:8080/auth/",username="admin",password="keycloakkeycloak",realm_name="master",verify=False)

    """ USERS """

    def get_user_id(self, username : str) -> str:
        self.connect()
        uid : str = self.keycloak_admin.get_user_id(username)
        return uid

    def get_users(self) -> Iterable[Dict[str, Any]]:
        # https://www.keycloak.org/docs-api/18.0/rest-api/index.html#_userrepresentation
        self.connect()
        o : Iterable[Dict[str, Any]] = self.keycloak_admin.get_users({})
        return o

    # TODO: what is this actually doing?
    def get_users_with_groups_and_roles(self) -> List[DDUser]:
        q = """select u.id, u.username, u.email, u.first_name, u.last_name, u.realm_id, u.enabled, ua.value as quota
        ,json_agg(g."id") as group, json_agg(g_parent."name") as group_parent1, json_agg(g_parent2."name") as group_parent2
        ,json_agg(r.name) as role
        from user_entity as u
        left join user_attribute as ua on ua.user_id=u.id and ua.name = 'quota'
        left join user_group_membership as ugm on ugm.user_id = u.id
        left join keycloak_group as g on g.id = ugm.group_id
        left join keycloak_group as g_parent on g.parent_group = g_parent.id
        left join keycloak_group as g_parent2 on g_parent.parent_group = g_parent2.id 
        left join user_role_mapping as rm on rm.user_id = u.id
        left join keycloak_role as r on r.id = rm.role_id
        group by u.id,u.username,u.email,u.first_name,u.last_name, u.realm_id, u.enabled, ua.value
        order by u.username"""

        (headers, users) = self.keycloak_pg.select_with_headers(q)

        users_with_lists = [
            list(l[:-4])
            + ([[]] if l[-4] == [None] else [list(set(l[-4]))])
            + ([[]] if l[-3] == [None] else [list(set(l[-3]))])
            + ([[]] if l[-3] == [None] else [list(set(l[-2]))])
            + ([[]] if l[-1] == [None] else [list(set(l[-1]))])
            for l in users
        ]

        users_with_lists = [
            list(l[:-4])
            + ([[]] if l[-4] == [None] else [list(set(l[-4]))])
            + ([[]] if l[-3] == [None] else [list(set(l[-3]))])
            + ([[]] if l[-3] == [None] else [list(set(l[-2]))])
            + ([[]] if l[-1] == [None] else [list(set(l[-1]))])
            for l in users_with_lists
        ]

        list_dict_users = [dict(zip(headers, r)) for r in users_with_lists]

        return list_dict_users

    def getparent(self, group_id : str, data : Iterable[Any]) -> str:
        # Recursively get full path from any group_id in the tree
        path = ""
        for item in data:
            if group_id == item[0]:
                path = self.getparent(item[2], data)
                path = f"{path}/{item[1]}"
        return path

    def get_group_path(self, group_id : str) -> str:
        # Get full path using getparent recursive func
        # RETURNS: String with full path
        q = """SELECT * FROM keycloak_group"""
        groups = self.keycloak_pg.select(q)
        return self.getparent(group_id, groups)

    def get_user_groups_paths(self, user_id : str) -> List[str]:
        # Get full paths for user grups
        # RETURNS list of paths
        q = """SELECT group_id FROM user_group_membership WHERE user_id = '%s'""" % (
            user_id
        )
        user_group_ids = self.keycloak_pg.select(q)

        paths = []
        for g in user_group_ids:
            paths.append(self.get_group_path(g[0]))
        return paths

    ## Too slow. Used the direct postgres
    # def get_users_with_groups_and_roles(self):
    #     self.connect()
    #     users=self.keycloak_admin.get_users({})
    #     for user in users:
    #         user['groups']=[g['path'] for g in self.keycloak_admin.get_user_groups(user_id=user['id'])]
    #         user['roles']= [r['name'] for r in self.keycloak_admin.get_realm_roles_of_user(user_id=user['id'])]
    #     return users

    def add_user(
        self,
        username : str,
        first : str,
        last : str,
        email : str,
        password : str,
        group : Any=False,
        password_temporary : bool=True,
        enabled : bool=True,
    ) -> Any:
        # RETURNS string with keycloak user id (the main id in this app)
        self.connect()
        username = username.lower()
        try:
            uid : Any = self.keycloak_admin.create_user(
                {
                    "email": email,
                    "username": username,
                    "enabled": enabled,
                    "firstName": first,
                    "lastName": last,
                    "credentials": [
                        {
                            "type": "password",
                            "value": password,
                            "temporary": password_temporary,
                        }
                    ],
                }
            )
        except Exception as e:
            log.error(traceback.format_exc())
            raise Error(
                "conflict",
                "user/email already exists: " + str(username) + "/" + str(email),
            )

        if group:
            path = "/" + group if group[1:] != "/" else group
            try:
                gid = self.keycloak_admin.get_group_by_path(
                    path=path, search_in_subgroups=False
                )["id"]
            except:
                self.keycloak_admin.create_group({"name": group})
                gid = self.keycloak_admin.get_group_by_path(path)["id"]
            self.keycloak_admin.group_user_add(uid, gid)
        return uid

    def update_user_pwd(self, user_id : str, password : str, password_temporary : bool=True) -> Any:
        # Updates
        payload = {
            "credentials": [
                {"type": "password", "value": password, "temporary": password_temporary}
            ]
        }
        self.connect()
        return self.keycloak_admin.update_user(user_id, payload)

    def user_update(self, user_id : str, enabled : bool, email : str, first : str, last : str, groups : Iterable[str]=[], roles : Iterable[str]=[]) -> Any:
        ## NOTE: Roles didn't seem to be updated/added. Also not confident with groups
        # Updates
        payload = {
            "enabled": enabled,
            "email": email,
            "firstName": first,
            "lastName": last,
            "groups": groups,
            "realmRoles": roles,
        }
        self.connect()
        return self.keycloak_admin.update_user(user_id, payload)

    def user_enable(self, user_id : str) -> Any:
        payload = {"enabled": True}
        self.connect()
        return self.keycloak_admin.update_user(user_id, payload)

    def user_disable(self, user_id : str) -> Any:
        payload = {"enabled": False}
        self.connect()
        return self.keycloak_admin.update_user(user_id, payload)

    def group_user_remove(self, user_id : str, group_id : str) -> Any:
        self.connect()
        return self.keycloak_admin.group_user_remove(user_id, group_id)

    # def add_user_role(self,user_id,role_id):
    #     self.connect()
    #     return self.keycloak_admin.assign_role(client_id=client_id, user_id=user_id, role_id=role_id, role_name="test")

    def remove_user_realm_roles(self, user_id : str, roles : Iterable[str]) -> Any:
        self.connect()
        roles = [
            r
            for r in self.get_user_realm_roles(user_id)
            if r["name"] in ["admin", "manager", "teacher", "student"]
        ]
        return self.keycloak_admin.delete_user_realm_role(user_id, roles)

    def delete_user(self, userid : str) -> Any:
        self.connect()
        return self.keycloak_admin.delete_user(user_id=userid)

    def get_user_groups(self, userid : str) -> Any:
        self.connect()
        return self.keycloak_admin.get_user_groups(user_id=userid)

    def get_user_realm_roles(self, userid : str) -> Any:
        self.connect()
        return self.keycloak_admin.get_realm_roles_of_user(user_id=userid)

    def add_user_client_role(self, client_id : str, user_id : str, role_id : str, role_name : str) -> Any:
        self.connect()
        return self.keycloak_admin.assign_client_role(
            client_id=client_id, user_id=user_id, role_id=role_id, role_name="test"
        )

    ## GROUPS
    def get_all_groups(self) -> Iterable[Any]:
        ## RETURNS ONLY MAIN GROUPS WITH NESTED subGroups list
        self.connect()
        return cast(Iterable[Any], self.keycloak_admin.get_groups())

    def get_groups(self, with_subgroups : bool=True) -> Iterable[Any]:
        ## RETURNS ALL GROUPS in root list
        self.connect()
        groups = self.keycloak_admin.get_groups()
        return get_recursive_groups(groups, [])

    def get_group_by_id(self, group_id : str) -> Any:
        self.connect()
        return self.keycloak_admin.get_group(group_id=group_id)

    def get_group_by_path(self, path : str, recursive : bool=True) -> Any:
        self.connect()
        return self.keycloak_admin.get_group_by_path(
            path=path, search_in_subgroups=recursive
        )

    def add_group(self, name : str, parent : str="", skip_exists : bool=False) -> Any:
        self.connect()
        parentId : Optional[str] = None
        if parent:
            parentId = self.get_group_by_path(parent)["id"]
        return self.keycloak_admin.create_group({"name": name}, parent=parentId)

    def delete_group(self, group_id : str) -> Any:
        self.connect()
        return self.keycloak_admin.delete_group(group_id=group_id)

    def group_user_add(self, user_id : str, group_id : str) -> Any:
        self.connect()
        return self.keycloak_admin.group_user_add(user_id, group_id)

    def add_group_tree(self, path : str) -> None:
        paths = kpath2kpaths(path)
        parent = "/"
        for path in paths:
            try:
                parent_path = "" if parent == "/" else parent
                # print("parent: "+str(parent_path)+" path: "+path.split("/")[-1])
                self.add_group(path.split("/")[-1], parent_path, skip_exists=True)
                parent = path
            except:
                # print(traceback.format_exc())
                log.warning("KEYCLOAK: Group :" + path + " already exists.")
                parent = path

    def add_user_with_groups_and_role(
        self, username : str, first : str, last : str, email : str, password : str, role : str, groups : Iterable[str]
    ) -> None:
        ## Add user
        uid = self.add_user(username, first, last, email, password)
        ## Add user to role
        log.info("User uid: " + str(uid) + " role: " + str(role))
        try:
            therole = role[0]
        except:
            therole = ""
        log.info(self.assign_realm_roles(uid, role))
        ## Create groups in user
        for g in groups:
            log.warning("Creating keycloak group: " + g)
            parts = g.split("/")
            parent_path = ""
            for i in range(1, len(parts)):
                # parent_id=None if parent_path==None else self.get_group(parent_path)['id']
                try:
                    self.add_group(parts[i], parent_path, skip_exists=True)
                except:
                    log.warning(
                        "Group "
                        + str(parent_path)
                        + " already exists. Skipping creation"
                    )
                    pass
                thepath = parent_path + "/" + parts[i]
                if thepath == "/":
                    log.warning(
                        "Not adding the user "
                        + username
                        + " to any group as does not have any..."
                    )
                    continue
                gid = self.get_group_by_path(path=thepath)["id"]

                log.warning(
                    "Adding "
                    + username
                    + " with uuid: "
                    + uid
                    + " to group "
                    + g
                    + " with uuid: "
                    + gid
                )
                self.keycloak_admin.group_user_add(uid, gid)

                parent_path += "/" + parts[i]

                # self.group_user_add(uid,gid)

    ## ROLES
    def get_roles(self) -> Iterable[Any]:
        self.connect()
        return cast(Iterable[Any], self.keycloak_admin.get_realm_roles())

    def get_role(self, name : str) -> Any:
        self.connect()
        return self.keycloak_admin.get_realm_role(name)

    def add_role(self, name : str, description : str="") -> Any:
        self.connect()
        return self.keycloak_admin.create_realm_role(
            {"name": name, "description": description}
        )

    def delete_role(self, name : str) -> Any:
        self.connect()
        return self.keycloak_admin.delete_realm_role(name)

    ## CLIENTS

    def get_client_roles(self, client_id : str) -> Any:
        self.connect()
        return self.keycloak_admin.get_client_roles(client_id=client_id)

    def add_client_role(self, client_id : str, name : str, description : str="") -> Any:
        self.connect()
        return self.keycloak_admin.create_client_role(
            client_id, {"name": name, "description": description, "clientRole": True}
        )

    ## SYSTEM
    def get_server_info(self) -> Any:
        self.connect()
        return self.keycloak_admin.get_server_info()

    def get_server_clients(self) -> Any:
        self.connect()
        return self.keycloak_admin.get_clients()

    def get_server_rsa_key(self) -> Any:
        self.connect()
        rsa_key = [
            k for k in self.keycloak_admin.get_keys()["keys"] if k["type"] == "RSA"
        ][0]
        return {"name": rsa_key["kid"], "certificate": rsa_key["certificate"]}

    ## REALM
    def assign_realm_roles(self, user_id : str, role : str) -> Any:
        self.connect()
        try:
            kcroles = [
                r for r in self.keycloak_admin.get_realm_roles() if r["name"] == role
            ]
        except:
            return False
        return self.keycloak_admin.assign_realm_roles(user_id=user_id, roles=kcroles)

    ## CLIENTS
    def delete_client(self, clientid : str) -> Any:
        self.connect()
        return self.keycloak_admin.delete_client(clientid)

    def add_client(self, client : str) -> Any:
        self.connect()
        return self.keycloak_admin.create_client(client)
