from __future__ import annotations

import os
from collections.abc import Callable
from datetime import datetime, timedelta
from decimal import Decimal
from pathlib import Path

import pandas as pd
from dask.distributed import Client as DaskClient
from distributed.utils_test import (
    client_no_amm,
)

from generalresearch.incite.base import GRLDatasets
from generalresearch.incite.collections.thl_web import (
    SessionDFCollection,
    WallDFCollection,
)
from generalresearch.incite.mergers.foundations.enriched_session import (
    EnrichedSessionMerge,
)
from generalresearch.incite.mergers.foundations.enriched_wall import EnrichedWallMerge
from generalresearch.managers.gr.team import MembershipManager, TeamManager
from generalresearch.models.gr.authentication import GRUser
from generalresearch.models.gr.business import Business
from generalresearch.models.gr.team import Membership, Team
from generalresearch.models.thl.product import Product
from generalresearch.models.thl.session import Session
from generalresearch.models.thl.user import User
from generalresearch.pg_helper import PostgresConfig
from generalresearch.redis_helper import RedisConfig


class TestTeam:

    def test_init(self, team: Team):

        assert isinstance(team, Team)
        assert isinstance(team.id, int)
        assert isinstance(team.uuid, str)

    def test_memberships_none(self, team: Team, gr_db: PostgresConfig):
        assert team.memberships is None

        team.prefetch_memberships(pg_config=gr_db)
        assert isinstance(team.memberships, list)
        assert len(team.memberships) == 0

    def test_memberships(
        self,
        team: Team,
        gr_user: GRUser,
        gr_user_factory: Callable[..., GRUser],
        membership_manager: MembershipManager,
        gr_db: PostgresConfig,
    ):
        assert team.memberships is None

        team.prefetch_memberships(pg_config=gr_db)
        assert isinstance(team.memberships, list)
        assert len(team.memberships) == 1
        assert team.memberships[0].user_id == gr_user.id

        # Create another new Membership
        membership_manager.create(team=team, gr_user=gr_user_factory())
        assert len(team.memberships) == 1
        team.prefetch_memberships(pg_config=gr_db)
        assert len(team.memberships) == 2

    def test_gr_users(
        self,
        team: Team,
        gr_user_factory: Callable[..., GRUser],
        membership_manager: MembershipManager,
        gr_db: PostgresConfig,
        gr_redis_config: RedisConfig,
    ):
        assert team.gr_users is None

        team.prefetch_gr_users(pg_config=gr_db, redis_config=gr_redis_config)
        assert isinstance(team.gr_users, list)
        assert len(team.gr_users) == 0

        # Create a new Membership
        membership_manager.create(team=team, gr_user=gr_user_factory())
        assert len(team.gr_users) == 0
        team.prefetch_gr_users(pg_config=gr_db, redis_config=gr_redis_config)
        assert len(team.gr_users) == 1

        # Create another Membership
        membership_manager.create(team=team, gr_user=gr_user_factory())
        assert len(team.gr_users) == 1
        team.prefetch_gr_users(pg_config=gr_db, redis_config=gr_redis_config)
        assert len(team.gr_users) == 2

    def test_businesses(
        self,
        team: Team,
        business: business: Business,
        team_manager: TeamManager,
        gr_db: PostgresConfig,
        gr_redis_config: RedisConfig,
    ):

        assert team.businesses is None

        team.prefetch_businesses(pg_config=gr_db, redis_config=gr_redis_config)
        assert isinstance(team.businesses, list)
        assert len(team.businesses) == 0

        team_manager.add_business(team=team, business=business)
        assert len(team.businesses) == 0
        team.prefetch_businesses(pg_config=gr_db, redis_config=gr_redis_config)
        assert len(team.businesses) == 1
        assert isinstance(team.businesses[0], Business)
        assert team.businesses[0].uuid == business.uuid

    def test_products(
        self,
        team: Team,
        product_factory: Callable[..., Product],
        thl_web_rr: PostgresConfig,
    ):

        assert team.products is None

        team.prefetch_products(thl_pg_config=thl_web_rr)
        assert isinstance(team.products, list)
        assert len(team.products) == 0

        product_factory(team=team)
        assert len(team.products) == 0
        team.prefetch_products(thl_pg_config=thl_web_rr)
        assert len(team.products) == 1
        assert isinstance(team.products[0], Product)


class TestTeamMethods:

    def test_cache_key(self, team: Team):
        assert isinstance(team.cache_key, str)
        assert ":" in team.cache_key
        assert str(team.uuid) in team.cache_key

    def test_set_cache(
        self,
        team: Team,
        gr_redis: RedisConfig,
        gr_db: PostgresConfig,
        thl_web_rr: PostgresConfig,
        gr_redis_config: RedisConfig,
        client_no_amm: DaskClient,
        mnt_filepath: GRLDatasets,
        mnt_gr_api_dir: Path,
        enriched_wall_merge: EnrichedWallMerge,
        enriched_session_merge: EnrichedSessionMerge,
    ):
        assert gr_redis.get(name=team.cache_key) is None

        team.set_cache(
            pg_config=gr_db,
            thl_web_rr=thl_web_rr: PostgresConfig,
            redis_config=gr_redis_config,
            client=client_no_amm,
            ds=mnt_filepath,
            mnt_gr_api=mnt_gr_api_dir,
            enriched_wall=enriched_wall_merge,
            enriched_session=enriched_session_merge,
        )

        assert gr_redis.hgetall(name=team.cache_key) is not None

    def test_set_cache_team(
        self,
        gr_user: GRUser,
        gr_db: PostgresConfig,
        thl_web_rr: PostgresConfig,
        product_factory: Callable[..., Product],
        team: Team,
        membership_factory: Callable[..., Membership],
        gr_redis_config: RedisConfig,
        mnt_filepath: GRLDatasets,
        mnt_gr_api_dir: Path,
        enriched_wall_merge: EnrichedWallMerge,
        enriched_session_merge: EnrichedSessionMerge,
    ):
        from generalresearch.models.gr.team import Team

        p1 = product_factory(team=team)
        membership_factory(team=team, gr_user=gr_user)

        team.set_cache(
            pg_config=gr_db,
            thl_web_rr=thl_web_rr: PostgresConfig,
            redis_config=gr_redis_config,
            client=client_no_amm,
            ds=mnt_filepath,
            mnt_gr_api=mnt_gr_api_dir,
            enriched_wall=enriched_wall_merge,
            enriched_session=enriched_session_merge,
        )

        team2 = Team.from_redis(
            uuid=team.uuid,
            fields=["id", "memberships", "gr_users", "businesses", "products"],
            gr_redis_config=gr_redis_config,
        )

        assert isinstance(team2, Team)
        assert isinstance(team2.products, list)
        assert isinstance(team2.gr_users, list)
        assert team.model_dump_json() == team2.model_dump_json()
        assert p1.uuid in [p.uuid for p in team2.products]
        assert len(team2.gr_users) == 1
        assert gr_user.id in [gru.id for gru in team2.gr_users]

    def test_prebuild_enriched_session_parquet(
        self,
        enriched_session_merge: EnrichedSessionMerge,
        client_no_amm: DaskClient,
        wall_collection: WallDFCollection,
        session_collection: SessionDFCollection,
        thl_web_rr: PostgresConfig,
        user_factory: Callable[..., User],
        start: datetime,
        session_factory: Callable[..., Session],
        product_factory: Callable[..., Product],
        delete_df_collection: Callable[..., None],
        mnt_filepath: GRLDatasets,
        mnt_gr_api_dir: Path,
        team: Team,
    ):

        delete_df_collection(coll=wall_collection)
        delete_df_collection(coll=session_collection)

        p1 = product_factory(team=team)
        p2 = product_factory(team=team)

        for p in [p1, p2]:
            u = user_factory(product=p)
            for i in range(50):
                session_factory(
                    user=u,
                    wall_count=1,
                    wall_req_cpi=Decimal("1.00"),
                    started=start + timedelta(minutes=i, seconds=1),
                )
        wall_collection.initial_load(client=None, sync=True)
        session_collection.initial_load(client=None, sync=True)

        enriched_session_merge.build(
            client=client_no_amm,
            session_coll=session_collection,
            wall_coll=wall_collection,
            pg_config=thl_web_rr: PostgresConfig,
        )

        team.prebuild_enriched_session_parquet(
            thl_pg_config=thl_web_rr: PostgresConfig,
            ds=mnt_filepath,
            client=client_no_amm,
            mnt_gr_api=mnt_gr_api_dir,
            enriched_session=enriched_session_merge,
        )

        # Now try to read from path
        df = pd.read_parquet(
            os.path.join(mnt_gr_api_dir, "pop_session", f"{team.file_key}.parquet")
        )
        assert isinstance(df, pd.DataFrame)

    def test_prebuild_enriched_wall_parquet(
        self,
        enriched_wall_merge: EnrichedWallMerge,
        client_no_amm: DaskClient,
        wall_collection: WallDFCollection,
        session_collection: EnrichedSessionMerge,
        thl_web_rr: PostgresConfig,
        user_factory: Callable[..., User],
        start: datetime,
        session_factory: Callable[..., Session],
        product_factory: Callable[..., Product],
        delete_df_collection: Callable[..., None],
        mnt_filepath: GRLDatasets,
        mnt_gr_api_dir: Path,
        team: Team,
    ):

        delete_df_collection(coll=wall_collection)
        delete_df_collection(coll=session_collection)

        p1 = product_factory(team=team)
        p2 = product_factory(team=team)

        for p in [p1, p2]:
            u = user_factory(product=p)
            for i in range(50):
                session_factory(
                    user=u,
                    wall_count=1,
                    wall_req_cpi=Decimal("1.00"),
                    started=start + timedelta(minutes=i, seconds=1),
                )
        wall_collection.initial_load(client=None, sync=True)
        session_collection.initial_load(client=None, sync=True)

        enriched_wall_merge.build(
            client=client_no_amm,
            session_coll=session_collection,
            wall_coll=wall_collection,
            pg_config=thl_web_rr: PostgresConfig,
        )

        team.prebuild_enriched_wall_parquet(
            thl_pg_config=thl_web_rr: PostgresConfig,
            ds=mnt_filepath,
            client=client_no_amm,
            mnt_gr_api=mnt_gr_api_dir,
            enriched_wall=enriched_wall_merge,
        )

        # Now try to read from path
        df = pd.read_parquet(
            os.path.join(mnt_gr_api_dir, "pop_event", f"{team.file_key}.parquet")
        )
        assert isinstance(df, pd.DataFrame)