from datetime import date
from functools import partial
from sqlalchemy import distinct, func, case, and_, or_, not_
from sqlalchemy.dialects.postgresql import DATERANGE, aggregate_order_by
from sqlalchemy.orm import selectinload
from typing import List, Any, Optional

from models.campaign import CampaignSourceType
from db import db
from sqlalchemy.orm.query import Query as SAQuery
from models.project_entity_type import ProjectEntityType, ProjectEntityAddType
from models import (
    Campaign,
    Project,
    Artist,
    Playlist,
    ProjectTargetItem,
    Polymorphable,
    user_label,
    UserProject,
    User,
    ArtistTeam,
    ArtistTeamUser,
    ProjectCampaign,
)
from constants.project_statuses import ProjectStatuses
from models.projects import ProjectCampaignStatus
from shared.queries import Queries as SharedQueries
from utils.string_utils import random_string
from utils import sorting_utils


Daterange = partial(func.daterange, type_=DATERANGE)


class DynamicFields:
    TOTAL_SPEND = 1
    CAMPAIGNS_COUNT = 2
    START_DATE = 3
    STATUS = 4
    USER_ROLE = 5
    SHARED_STATUS = 6
    TARGET_ITEMS = 7


class BuilderJoin:
    model_class: db.Model
    statement: Any
    is_inner: bool = True

    def __init__(self, model_class, statement, inner):
        self.model_class = model_class
        self.statement = statement
        self.is_inner = inner


class ProjectsQueryBuilder:
    projects_sort_map = {
        "id": (Project.id,),
        "project": (Project.name, Project.id),
        "plannedBudget": (Project.budget, Project.name, Project.id),
        "budget": (Project.budget, Project.name, Project.id),
        "allocation": (Project.allocation, Project.name, Project.id),
        "totalSpend": ("total_spend", Project.name, Project.id),
        "earliestStartDate": ("start_date", Project.name, Project.id),
        "earliestEndDate": (Project.end_date, Project.name, Project.id),
        "status": ("status", "start_date", Project.name, Project.id),
        "userRoles": ("user_roles", Project.name, Project.id),
        "artistName": ("artists_names", Project.name, Project.id),
        "playlistName": ("playlists_names", Project.name, Project.id),
        "created_at": (Project.created_at,)
    }

    def __init__(self, user_id: Optional[int] = None):
        # This token used to make subqueries unique for the case when multiple requests are executed in one
        # transaction
        self.token = random_string(5)

        self.user_id = user_id
        self.project_targets_data_query = self.__get_project_targets_data_query().cte(f"project_targets{self.token}")
        self.project_metadata_query = self.__get_projects_metadata_query().cte(f"projects_meta{self.token}")
        self.project_campaigns_data_query = self.__get_campaigns_data_query().cte(f"campaign_metadata{self.token}")
        self.artist_team_query = self.__get_artist_team_data_query().cte(f"artist_team_data{self.token}")
        self.artist_roles_query = self.__get_artist_team_roles_for_project_query(self.user_id).cte(
            f"artist_roles{self.token}"
        )
        self.user_roles_data_query = self.__get_roles_for_project_query(self.user_id).cte(f"user_roles{self.token}")
        self.filter_items = [Project.is_deleted.is_(False)]
        self.sort_expression = sorting_utils.sorting_expression(self.projects_sort_map, "id")
        self.dynamic_fields = [
            DynamicFields.TOTAL_SPEND,
            DynamicFields.CAMPAIGNS_COUNT,
            DynamicFields.START_DATE,
            DynamicFields.STATUS,
            DynamicFields.USER_ROLE,
            DynamicFields.SHARED_STATUS,
            DynamicFields.TARGET_ITEMS,
        ]
        self.additional_joins: List[BuilderJoin] = []
        self.limit = 20
        self.offset = 0

    def __get_project_targets_data_query(self) -> SAQuery:
        primary_artist_or_playlist = [ProjectEntityType.PRIMARY_ARTIST.value, ProjectEntityType.PLAYLIST.value]
        return (
            db.session.query(
                Project.id.label("project_id"),
                func.STRING_AGG(Artist.name, "").label("artists_names"),
                func.STRING_AGG(Playlist.name, "").label("playlists_names"),
            )
            .join(
                ProjectTargetItem,
                and_(
                    Project.id == ProjectTargetItem.project_id,
                    ProjectTargetItem.entity_type.in_(primary_artist_or_playlist),
                ),
            )
            .join(Polymorphable, ProjectTargetItem.entity_id == Polymorphable.id)
            .outerjoin(Artist, Polymorphable.id == Artist.id)
            .outerjoin(Playlist, Polymorphable.id == Playlist.id)
            .group_by(Project.id)
        )

    def __get_campaigns_data_query(self) -> SAQuery:
        internal_campaign = case(
            [
                (
                    or_(Campaign.source.is_(None), Campaign.source.in_(CampaignSourceType.all_internal_sources())),
                    Campaign.id
                )
            ]
        )
        external_campaign = case([(Campaign.source.in_(CampaignSourceType.all_external_sources()), Campaign.id)])
        pending_campaign = case([(ProjectCampaign.status == ProjectCampaignStatus.PENDING.value, ProjectCampaign.id)])

        return (
            db.session.query(
                Project.id.label("project_id"),
                func.count(distinct(Campaign.id)).label("total_campaigns_count"),
                func.coalesce(func.count(distinct(internal_campaign)), 0).label("internal"),
                func.coalesce(func.count(distinct(external_campaign)), 0).label("external"),
                func.coalesce(func.count(distinct(pending_campaign)), 0).label("pending"),
                func.sum(Campaign.planned_budget).label("total_spend"),
            )
            .select_from(Project)
            .outerjoin(
                Campaign,
                and_(
                    Campaign.project_id == Project.id,
                    Campaign.is_deleted.is_(False)
                )
            )
            .outerjoin(ProjectCampaign, ProjectCampaign.project_id == Project.id)
            .group_by(Project.id)
        )

    def __get_artist_team_data_query(self) -> SAQuery:
        return (
            db.session.query(ArtistTeam.id, ArtistTeam.artist_id, ArtistTeam.label_id, ProjectTargetItem.project_id)
            .join(
                ProjectTargetItem,
                and_(
                    ProjectTargetItem.is_deleted.is_(False),
                    ProjectTargetItem.entity_type == ProjectEntityType.PRIMARY_ARTIST.value,
                    ProjectTargetItem.add_type == ProjectEntityAddType.LOCKED.value,
                    ProjectTargetItem.entity_id == ArtistTeam.artist_id,
                ),
            )
            .join(Polymorphable, Polymorphable.id == ProjectTargetItem.entity_id)
            .join(Artist, and_(Artist.id == Polymorphable.id, Artist.is_unknown.is_(False)))
        )

    def __get_roles_for_project_query(self, user_id) -> SAQuery:
        return (
            db.session.query(
                Project.id.label("project_id"),
                func.array_agg(aggregate_order_by(UserProject.role, UserProject.role.asc())).label("user_project_roles")
            )
            .select_from(Project)
            .join(UserProject,
                  and_(
                      UserProject.project_id == Project.id,
                      UserProject.user_id == self.user_id
                  ))
            .group_by(Project.id)
        )

    def __get_artist_team_roles_for_project_query(self, user_id) -> SAQuery:
        return (
            db.session.query(
                ArtistTeam.id.label("artist_team_id"),
                func.array_agg(
                    aggregate_order_by(ArtistTeamUser.role, ArtistTeamUser.role.asc())
                ).label("artist_team_roles")
            )
            .select_from(ArtistTeam)
            .join(
                ArtistTeamUser,
                and_(
                    ArtistTeamUser.artist_team_id == ArtistTeam.id,
                    ArtistTeamUser.user_id == self.user_id
                )
            )
            .group_by(ArtistTeam.id)
        )

    def __get_projects_metadata_query(self) -> SAQuery:
        return (
            db.session.query(
                Project.id.label("project_id"),
                Project.initial_start_date.label("start_date"),
                case(
                    [
                        (func.max(Project.end_date) < func.current_date(), ProjectStatuses.COMPLETED.value),
                        (Project.initial_start_date <= func.current_date(), ProjectStatuses.INPROGRESS.value),
                    ],
                    else_=ProjectStatuses.SCHEDULED.value,
                ).label("status"),
            )
            .select_from(Project)
            .group_by(Project.id)
        )

    def set_dynamic_fields(self, fields: List[DynamicFields]):
        self.dynamic_fields = fields

    def __base_fields_set(self):
        fields_set = [Project]
        if DynamicFields.TOTAL_SPEND in self.dynamic_fields:
            fields_set.append(func.coalesce(self.project_campaigns_data_query.c.total_spend, 0).label("total_spend"))
        if DynamicFields.CAMPAIGNS_COUNT in self.dynamic_fields:
            fields_set.append(
                func.coalesce(self.project_campaigns_data_query.c.total_campaigns_count, 0).label("campaigns_count")
            )
            fields_set.append(
                func.coalesce(self.project_campaigns_data_query.c.internal, 0).label("internal_campaigns_count")
            )
            fields_set.append(
                func.coalesce(self.project_campaigns_data_query.c.external, 0).label("external_campaigns_count")
            )
            fields_set.append(
                func.coalesce(self.project_campaigns_data_query.c.pending, 0).label("pending_campaigns_count")
            )
        if DynamicFields.START_DATE in self.dynamic_fields:
            fields_set.append(self.project_metadata_query.c.start_date.label("start_date"))
        if DynamicFields.STATUS in self.dynamic_fields:
            fields_set.append(
                func.coalesce(self.project_metadata_query.c.status, ProjectStatuses.SCHEDULED.value).label("status")
            )
        if DynamicFields.USER_ROLE in self.dynamic_fields:
            fields_set.append(
                SharedQueries.roles_for_project_query(
                    self.user_id,
                    self.user_roles_data_query.c.user_project_roles,
                    self.artist_roles_query.c.artist_team_roles
                ).label("user_roles"))
        if DynamicFields.SHARED_STATUS in self.dynamic_fields:
            fields_set.append(SharedQueries.project_shared_status_query(self.user_id).label("shared_status"))
        if DynamicFields.TARGET_ITEMS in self.dynamic_fields:
            fields_set.append(self.project_targets_data_query.c.artists_names)
            fields_set.append(self.project_targets_data_query.c.playlists_names)

        return fields_set

    def base_query(self, fields_set):
        query = (
            db.session.query(*fields_set)
            .join(
                ProjectTargetItem,
                and_(ProjectTargetItem.project_id == Project.id, ProjectTargetItem.is_deleted.is_(False)),
            )
            .join(Polymorphable, Polymorphable.id == ProjectTargetItem.entity_id)
            .outerjoin(
                Artist,
                and_(
                    Artist.id == Polymorphable.id,
                    ProjectTargetItem.entity_type == ProjectEntityType.PRIMARY_ARTIST.value,
                ),
            )
        )
        if self.user_id is not None:
            query = (
                query.join(User, User.id == self.user_id)
                .outerjoin(
                    UserProject,
                    and_(
                        UserProject.project_id == Project.id,
                        UserProject.user_id == User.id
                    ))
                .outerjoin(
                    user_label, and_(user_label.c.label_id == Project.label_id, user_label.c.user_id == self.user_id)
                )
                .outerjoin(
                    self.user_roles_data_query, Project.id == self.user_roles_data_query.c.project_id
                )
                .outerjoin(
                    self.artist_team_query,
                    and_(
                        self.artist_team_query.c.label_id == Project.label_id,
                        self.artist_team_query.c.project_id == Project.id,
                    ),
                )
                .outerjoin(
                    ArtistTeamUser,
                    and_(
                        self.artist_team_query.c.id == ArtistTeamUser.artist_team_id,
                        ArtistTeamUser.user_id == self.user_id,
                    ),
                )
                .outerjoin(
                    self.artist_roles_query, self.artist_team_query.c.id == self.artist_roles_query.c.artist_team_id
                )
            )
        if DynamicFields.START_DATE in self.dynamic_fields or DynamicFields.STATUS in self.dynamic_fields:
            query = query.outerjoin(self.project_metadata_query, Project.id == self.project_metadata_query.c.project_id)
        if DynamicFields.CAMPAIGNS_COUNT in self.dynamic_fields or DynamicFields.TOTAL_SPEND in self.dynamic_fields:
            query = query.outerjoin(
                self.project_campaigns_data_query, Project.id == self.project_campaigns_data_query.c.project_id
            )
        if DynamicFields.TARGET_ITEMS in self.dynamic_fields:
            query = query.outerjoin(
                self.project_targets_data_query, Project.id == self.project_targets_data_query.c.project_id
            )

        return query.distinct()

    def only_claimed(self):
        self.filter_items.append(Project.is_claimed.is_(True))
        return self

    def only_unclaimed(self):
        self.filter_items.append(Project.is_claimed.is_(False))
        return self

    def only_accessible_projects(self):
        is_labelmate = user_label.c.user_id.isnot(None)
        can_see_project = or_(
            and_(is_labelmate, Project.is_confidential.is_(False)).self_group(),
            and_(is_labelmate, User.is_admin).self_group(),
        )

        in_artist_team = and_(ArtistTeamUser.user_id == self.user_id, Project.is_confidential.is_(False))

        self.filter_items.append(or_(can_see_project, in_artist_team, UserProject.user_id.isnot(None)))
        return self

    def only_accessible_projects_by_artist(self):
        is_labelmate = user_label.c.user_id.isnot(None)
        can_see_project = or_(
            and_(is_labelmate, Project.is_confidential.is_(False), ArtistTeamUser.user_id == self.user_id).self_group(),
            and_(is_labelmate, User.is_admin).self_group(),
        )

        in_artist_team = and_(ArtistTeamUser.user_id == self.user_id, Project.is_confidential.is_(False))

        self.filter_items.append(or_(
            can_see_project,
            in_artist_team,
            and_(UserProject.user_id.isnot(None), ArtistTeamUser.user_id == self.user_id)
        ))
        return self

    def with_currency(self, currency: str):
        self.filter_items.append(Project.currency == currency)
        return self

    def filtered_by_search(self, query: str):
        found_ids = Project.search(expression=query)
        self.filter_items.append(Project.id.in_(found_ids))
        return self

    def filtered_by_type(self, type: str):
        self.filter_items.append(Polymorphable.type == type)
        return self

    def filter_by_statuses(self, statuses: List[int]):
        self.filter_items.append(
            func.coalesce(self.project_metadata_query.c.status, ProjectStatuses.SCHEDULED.value).in_(statuses)
        )
        return self

    def filter_by_status(self, status: int):
        self.filter_items.append(
            func.coalesce(self.project_metadata_query.c.status, ProjectStatuses.SCHEDULED.value) == status
        )
        return self

    def filtered_by_date_range(self, start_date: date, end_date: date):
        project_range = Daterange(self.project_metadata_query.c.start_date, Project.end_date, "[]")
        filter_range = Daterange(start_date, end_date, "[]")
        self.filter_items.append(project_range.overlaps(filter_range))
        return self

    def filtered_by_labels(self, labels: List[int]):
        self.filter_items.append(Project.label_id.in_(labels))
        return self

    def filtered_by_artist(self, artist_external_id: str, prs_only: bool = False):
        self.filter_items.append(Artist.external_id == artist_external_id)

        if prs_only:
            self.filter_items.append(ProjectTargetItem.add_type == ProjectEntityAddType.LOCKED.value)

        return self

    def only_mine(self):
        self.filter_items.append(
            or_(
                UserProject.id.isnot(None),
                and_(ArtistTeamUser.id.isnot(None), Project.is_confidential.is_(False))
            )
        )
        return self

    def sort_by(self, sort_field: str):
        self.sort_expression = sorting_utils.sorting_expression(self.projects_sort_map, sort_field)
        return self

    def limit_to(self, limit: int):
        self.limit = limit
        return self

    def offset_by(self, offset: int):
        self.offset = offset
        return self

    def items_query(self, use_pagination: Optional[bool] = True):
        query = self.base_query(self.__base_fields_set())

        for join in self.additional_joins:
            if join.is_inner:
                query = query.join(join.model_class, join.statement)
            else:
                query = query.outerjoin(join.model_class, join.statement)

        for f in self.filter_items:
            query = query.filter(f)

        query = query.options(
            selectinload(Project.target_items),
        )
        query = query.order_by(*self.sort_expression)
        if not use_pagination:
            return query
        return query.limit(self.limit).offset(self.offset)

    def count_query(self):
        query = self.base_query((func.count(distinct(Project.id)),))

        for f in self.filter_items:
            query = query.filter(f)

        return query