# By: Riasat Ullah
# This file contains functions for handling tags related database queries.
import json

from dbqueries import db_policies, db_services, db_teams, db_users
from psycopg2 import errorcodes
from utils import constants, errors, key_manager, var_names
from uuid import UUID
import configuration as configs
import datetime
import psycopg2


def create_tag(conn, timestamp, org_id, tag, description=None, tag_policies=None, tag_users=None, tag_teams=None,
               tag_services=None):
    '''
    Create a new tag.
    :param conn: db connection
    :param timestamp: timestamp when the request was made
    :param org_id: ID of the organization
    :param tag: the tag
    :param description: description of what the tag is for
    :param tag_policies: (list) of concealed policy ref IDs
    :param tag_users: (list) of preferred usernames
    :param tag_teams: (list) of concealed team ref IDs
    :param tag_services: (list) of concealed service ref IDs
    :errors: AssertionError, DatabaseError, LookupError
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)
    assert isinstance(tag, str)
    if description is not None:
        assert isinstance(description, str)

    tag = tag.lower()
    associations = map_tag_associations_from_ref_ids(
        conn, timestamp, org_id, tag_policies, tag_users, tag_teams, tag_services)

    query = "select create_tag(%s, %s, %s, %s, %s, %s);"
    query_params = (timestamp, constants.end_timestamp, org_id, tag, description, json.dumps(associations),)
    try:
        conn.execute(query, query_params)
    except psycopg2.IntegrityError as e:
        if e.pgcode == errorcodes.CHECK_VIOLATION:
            raise LookupError(errors.err_tag_duplicate)
        else:
            raise


def edit_tag(conn, timestamp, org_id, tag, description=None, tag_policies=None, tag_users=None, tag_teams=None,
             tag_services=None):
    '''
    Edit an existing tag.
    :param conn: db connection
    :param timestamp: timestamp when the request was made
    :param org_id: ID of the organization
    :param tag: the tag
    :param description: description of what the tag is for
    :param tag_policies: (list) of concealed policy ref IDs
    :param tag_users: (list) of preferred usernames
    :param tag_teams: (list) of concealed team ref IDs
    :param tag_services: (list) of concealed service ref IDs
    :return:
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)
    assert isinstance(tag, str)
    if description is not None:
        assert isinstance(description, str)

    tag = tag.lower()
    associations = map_tag_associations_from_ref_ids(
        conn, timestamp, org_id, tag_policies, tag_users, tag_teams, tag_services)

    query = "select edit_tag(%s, %s, %s, %s, %s, %s);"
    query_params = (timestamp, constants.end_timestamp, org_id, tag, description, json.dumps(associations),)
    try:
        conn.execute(query, query_params)
    except psycopg2.IntegrityError as e:
        if e.pgcode == errorcodes.CHECK_VIOLATION:
            raise LookupError(errors.err_unknown_resource)
        else:
            raise


def delete_tag(conn, timestamp, organization_id, tag):
    '''
    Deletes a tag.
    :param conn: db connection
    :param timestamp: timestamp when the request is being made
    :param organization_id: ID of the organization
    :param tag: the tag to delete
    :errors: AssertionError, DatabaseError
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(organization_id, int)
    assert isinstance(tag, str)

    query = '''
            begin;

            update tags set end_timestamp = %(timestamp)s
            where start_timestamp <= %(timestamp)s
                and end_timestamp > %(timestamp)s
                and organization_id = %(org_id)s
                and tag = %(tg)s;

            update tag_associations set end_timestamp = %(timestamp)s
            where start_timestamp <= %(timestamp)s
                and end_timestamp > %(timestamp)s
                and organization_id = %(org_id)s
                and tag = %(tg)s;

            update monitor_check_incident_tags set end_timestamp = %(timestamp)s
            where start_timestamp <= %(timestamp)s
                and end_timestamp > %(timestamp)s
                and check_id in (
                    select check_id from monitor_checks
                    where start_timestamp <= %(timestamp)s
                        and end_timestamp > %(timestamp)s
                        and organization_id = %(org_id)s
                )
                and tag = %(tg)s;

            end;
            '''
    query_params = {'timestamp': timestamp, 'org_id': organization_id, 'tg': tag}
    try:
        conn.execute(query, query_params)
    except psycopg2.DataError:
        raise


def save_component_tags(conn, timestamp, org_id, tags, policy_ref_id=None, preferred_username=None,
                        team_ref_id=None, service_ref_id=None):
    '''
    Save the tags of a component.
    :param conn: db connection
    :param timestamp: timestamp when the request is being made
    :param org_id: ID of the organization
    :param tags: (list) of tags; can also be null
    :param policy_ref_id: reference ID of the policy to add the tags to
    :param preferred_username: preferred username of the user to add the tags to
    :param team_ref_id: reference ID of the team to add the tags to
    :param service_ref_id: reference ID of the service to add the tags to
    :errors: AssertionError, DatabaseError
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)
    comp_type_id, comp_id = None, None
    if policy_ref_id is not None:
        comp_id = db_policies.list_policy_ids_from_ref_ids(conn, timestamp, org_id, [policy_ref_id], unmasked=False)[0]
        comp_type_id = configs.policy_component_type_id
    if preferred_username is not None:
        found_users = db_users.get_user_ids_from_preferred_usernames(conn, timestamp, org_id, [preferred_username])
        if len(found_users) == 0:
            raise LookupError(errors.err_tag_users)
        else:
            comp_id = found_users[0]
            comp_type_id = None
    if team_ref_id is not None:
        comp_id = db_teams.get_team_ids_from_ref_ids(conn, timestamp, org_id, [team_ref_id], unmasked=False,
                                                     as_dict=False)[0]
        comp_type_id = configs.team_component_type_id
    if service_ref_id is not None:
        comp_id = db_services.list_service_ids_from_ref_ids(conn, timestamp, org_id, [service_ref_id], unmasked=False,
                                                            as_dict=False)[0]
        comp_type_id = configs.service_component_type_id

    if comp_id is not None:
        query = " select save_component_tags(%s, %s, %s, %s::smallint, %s, %s); "
        query_params = (timestamp, constants.end_timestamp, org_id, comp_type_id, comp_id,
                        tags if tags is not None else None,)
        try:
            conn.execute(query, query_params)
        except psycopg2.DatabaseError:
            raise
    else:
        raise LookupError(errors.err_unknown_resource)


def list_tags(conn, timestamp, org_id, with_preferred_username=False):
    '''
    List all the tags in an organization and provide details of all the components they are associated with.
    :param conn: db connection
    :param timestamp: timestamp when this request is being made
    :param org_id: ID of the organization
    :param with_preferred_username: True if user associations should be mapped with preferred username;
                    False if it should be mapped with the user policy id
    :return: (list of dict) -> [{tag: ..., policies: ..., users: ..., teams: ..., services: ...}, ...]
    :errors: AssertionError, DatabaseError
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)

    query = '''
            with t1 as (
                select tag, description from tags
                where start_timestamp <= %(timestamp)s
                    and end_timestamp > %(timestamp)s
                    and organization_id = %(org_id)s
            )
            , t2 as (
                select taga.tag, json_agg(json_build_object(
                    'policy_ref_id', pol.policy_ref_id,
                    'policy_name', pol.policy_name
                )) as tag_policies
                from t1
                join tag_associations as taga using (tag)
                join policies as pol
                    on taga.component_type_id = %(pol_comp_type_id)s
                        and taga.component_id = pol.policyid
                        and pol.is_group_policy = true
                        and pol.start_timestamp <= %(timestamp)s
                        and pol.end_timestamp > %(timestamp)s
                        and pol.organization_id = %(org_id)s
                where taga.start_timestamp <= %(timestamp)s
                    and taga.end_timestamp > %(timestamp)s
                    and taga.organization_id = %(org_id)s
                group by taga.tag
            )
            , t3 as (
                select taga.tag, json_agg(json_build_object(
                    'preferred_username', users.preferred_username,
                    'display_name', users.first_name || ' ' || users.last_name,
                    'policy_ref_id', pol.policy_ref_id
                )) as tag_users
                from t1
                join tag_associations as taga using (tag)
                join users
                    on taga.component_type_id is null
                        and taga.component_id = users.user_id
                        and users.start_timestamp <= %(timestamp)s
                        and users.end_timestamp > %(timestamp)s
                        and users.organization_id = %(org_id)s
                left join policies as pol
                    on users.policyid = pol.policyid
                        and pol.start_timestamp <= %(timestamp)s
                        and pol.end_timestamp > %(timestamp)s
                        and pol.organization_id = %(org_id)s
                where taga.start_timestamp <= %(timestamp)s
                    and taga.end_timestamp > %(timestamp)s
                    and taga.organization_id = %(org_id)s
                group by taga.tag
            )
            , t4 as (
                select taga.tag, json_agg(json_build_object(
                    'team_ref_id', tm.team_ref_id,
                    'team_name', tm.team_name
                )) as tag_teams
                from t1
                join tag_associations as taga using (tag)
                join teams as tm
                    on taga.component_type_id = %(team_comp_type_id)s
                        and taga.component_id = tm.team_id
                        and tm.start_timestamp <= %(timestamp)s
                        and tm.end_timestamp > %(timestamp)s
                        and tm.organization_id = %(org_id)s
                where taga.start_timestamp <= %(timestamp)s
                    and taga.end_timestamp > %(timestamp)s
                    and taga.organization_id = %(org_id)s
                group by taga.tag
            )
            , t5 as (
                select taga.tag, json_agg(json_build_object(
                    'service_ref_id', srv.service_ref_id,
                    'service_name', srv.service_name
                )) as tag_services
                from t1
                join tag_associations as taga using (tag)
                join services as srv
                    on taga.component_type_id = %(srv_comp_type_id)s
                        and taga.component_id = srv.serviceid
                        and srv.start_timestamp <= %(timestamp)s
                        and srv.end_timestamp > %(timestamp)s
                        and srv.organization_id = %(org_id)s
                where taga.start_timestamp <= %(timestamp)s
                    and taga.end_timestamp > %(timestamp)s
                    and taga.organization_id = %(org_id)s
                group by taga.tag
            )
            select t1.tag, t1.description, t2.tag_policies, t3.tag_users, t4.tag_teams, t5.tag_services
            from t1
            left join t2 using(tag)
            left join t3 using(tag)
            left join t4 using(tag)
            left join t5 using(tag);
            '''
    query_params = {'timestamp': timestamp, 'org_id': org_id,
                    'pol_comp_type_id': configs.policy_component_type_id,
                    'team_comp_type_id': configs.team_component_type_id,
                    'srv_comp_type_id': configs.service_component_type_id}
    try:
        result = conn.fetch(query, query_params)
        data = []
        for tag, desc, pols, usrs, tms, srvs in result:
            pols = [[item[var_names.policy_name],
                     key_manager.conceal_reference_key(UUID(item[var_names.policy_ref_id]))]
                    for item in pols] if pols is not None else []

            if with_preferred_username:
                usrs = [[item[var_names.display_name], item[var_names.preferred_username]]
                        for item in usrs] if usrs is not None else []
            else:
                usrs = [[item[var_names.display_name],
                         key_manager.conceal_reference_key(UUID(item[var_names.policy_ref_id]))]
                        for item in usrs if item[var_names.policy_ref_id] is not None]\
                    if usrs is not None else []

            tms = [[item[var_names.team_name], key_manager.conceal_reference_key(UUID(item[var_names.team_ref_id]))]
                   for item in tms] if tms is not None else []

            srvs = [[item[var_names.service_name],
                     key_manager.conceal_reference_key(UUID(item[var_names.service_ref_id]))]
                    for item in srvs] if srvs is not None else []

            data.append({
                var_names.tag: tag,
                var_names.description: desc,
                var_names.policies: pols,
                var_names.users: usrs,
                var_names.teams: tms,
                var_names.services: srvs
            })
        return data
    except psycopg2.DatabaseError:
        raise


def get_basic_tags_list(conn, timestamp, org_id):
    '''
    Get the list of all tags of an organization.
    :param conn: db connection
    :param timestamp: timestamp when this request is being made
    :param org_id: ID of the organization
    :return: (list) of tags
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)

    query = '''
            select tag from tags
            where start_timestamp <= %s
                and end_timestamp > %s
                and organization_id = %s;
            '''
    query_params = (timestamp, timestamp, org_id,)
    try:
        result = conn.fetch(query, query_params)
        data = []
        for item in result:
            data.append(item[0])
        return data
    except psycopg2.DatabaseError:
        raise


def map_tag_associations_from_ref_ids(conn, timestamp, org_id, tag_policies=None, tag_users=None, tag_teams=None,
                                      tag_services=None):
    '''
    Get the internal IDs of the components associated to a tag using their reference IDs.
    In the case of users, their preferred username is the reference ID.
    :param conn: db connection
    :param timestamp: timestamp when the request was made
    :param org_id: ID of the organization
    :param tag_policies: (list) of concealed policy ref IDs
    :param tag_users: (list) of preferred usernames
    :param tag_teams: (list) of concealed team ref IDs
    :param tag_services: (list) of concealed service ref IDs
    :return: (list of dict) -> [{component_type_id: ..., component_id: ...}, ...]
    '''
    assert isinstance(timestamp, datetime.datetime)
    assert isinstance(org_id, int)

    associations = []
    if tag_policies is not None:
        pol_ids = db_policies.list_policy_ids_from_ref_ids(conn, timestamp, org_id, tag_policies, unmasked=False)
        for id_ in pol_ids:
            associations.append({var_names.component_type_id: configs.policy_component_type_id,
                                 var_names.component_id: id_})
    if tag_users is not None:
        user_ids = db_users.get_user_ids_from_preferred_usernames(conn, timestamp, org_id, tag_users)
        for id_ in user_ids:
            associations.append({var_names.component_type_id: None, var_names.component_id: id_})
    if tag_teams is not None:
        tm_ids = db_teams.get_team_ids_from_ref_ids(conn, timestamp, org_id, tag_teams, unmasked=False, as_dict=False)
        for id_ in tm_ids:
            associations.append({var_names.component_type_id: configs.team_component_type_id,
                                 var_names.component_id: id_})
    if tag_services is not None:
        srv_ids = db_services.list_service_ids_from_ref_ids(conn, timestamp, org_id, tag_services, unmasked=False,
                                                            as_dict=False)
        for id_ in srv_ids:
            associations.append({var_names.component_type_id: configs.service_component_type_id,
                                 var_names.component_id: id_})

    return associations