from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List

import polars as pl


class DateAgg(Enum):
    DAY = 'DAY'
    MONTH = 'MONTH'
    QUARTER = 'QUARTER'
    YEAR = 'YEAR'


@dataclass
class StatsModel:
    """Class that models the output of the Tumblr stats script."""
    # The blog in question.
    blog_name: str

    # Contains posts.
    posts: pl.DataFrame

    # Contains queued posts.
    queued_posts: pl.DataFrame

    # Contains drafted posts, both original and not original.
    drafts: pl.DataFrame

    # Contains activity (notifications).
    activity: pl.DataFrame 

    # Any tags used.
    tags: List[str] = field(default_factory=list)

    # Total count of posts processed.
    total_posts: int = field(init=False)

    # Total original posts (for blog_name) processed.
    total_original_posts: int = field(init=False)

    # Total original post (for blog_name) notes processed.
    total_original_post_notes: int = field(init=False)

    # Note stats for original posts by YYYY-MM-DD.
    original_post_notes_by_date: List[Dict[str, Any]] = field(init=False)

    # Note stats for original posts by month and year.
    original_post_notes_by_month_and_year:  List[Dict[str, Any]] = field(init=False)

    # Note stats for original posts by quarter and year.
    original_post_notes_by_qtr_and_year:  List[Dict[str, Any]] = field(init=False)

    # Note stats for original posts by year.
    original_post_notes_by_year:  List[Dict[str, Any]] = field(init=False)

    # Tags ranked from most popular to least popular by notes.
    most_popular_tags: List[Dict[str, Any]] = field(init=False)

    # Original posts ranked from most popular to least popular by notes.
    ranked_original_posts: List[Dict[str, Any]] = field(init=False)

    # Original post note stats by interaction month and year.
    original_post_notes_by_interaction_date: List[Dict[str, Any]] = field(init=False)

    # Original post note stats by interaction month and year.
    original_post_notes_by_interaction_month_and_year: List[Dict[str, Any]] = field(init=False)

    # Top blogs by original post note count.
    top_interactors: List[Dict[str, Any]] = field(init=False)

    # Queued posts (both original and not original), sorted in publish order.
    ordered_queue: List[Dict[str, Any]] = field(init=False)

    # Top tags in queued posts (both original and not original), sorted by post count.
    top_queued_tags: List[Dict[str, Any]] = field(init=False)

    # Top tags in drafts (both original and not original), sorted by post count.
    top_draft_tags: List[Dict[str, Any]] = field(init=False)

    def __post_init__(self):
        self.total_posts = self.calculate_total_posts()
        self.total_original_posts = self.calculate_total_original_posts()
        self.total_original_post_notes = self.calculate_original_post_notes()
        self.original_post_notes_by_date = self.collect_note_stats_by_original_post_year_month_day()
        self.original_post_notes_by_month_and_year = self.collect_note_stats_by_original_post_month_and_year()
        self.original_post_notes_by_qtr_and_year = self.collect_note_stats_by_original_post_qtr_and_year()
        self.original_post_notes_by_year = self.collect_note_stats_by_original_post_year()
        self.most_popular_tags = self.determine_most_popular_tags(
            self.get_original_posts(), "note_count")
        self.ranked_original_posts = self.rank_original_posts()
        self.original_post_notes_by_interaction_month_and_year = self.collect_note_stats_by_op_interaction_month_and_year()
        self.original_post_notes_by_interaction_date = self.collect_note_stats_by_op_interaction_year_month_day()
        self.top_interactors = self.determine_top_interactors()
        self.ordered_queue = self.determine_ordered_queue()
        self.top_queued_tags = self.determine_most_popular_tags(
            self.queued_posts, 'post_count')
        self.top_draft_tags = self.determine_most_popular_tags(
            self.drafts, 'post_count')

    @staticmethod
    def determine_most_popular_tags(df: pl.DataFrame, sort_key: str) -> List[Dict[str, Any]]:
        if df.is_empty():
            return []

        grouped_tags = df.explode(
            'tags', keep_nulls=False, empty_as_null=False).group_by('tags')
        result: pl.DataFrame = pl.DataFrame()
        by_post_count = grouped_tags.agg(pl.count().alias(
            'post_count')).sort('post_count', descending=True)
        if sort_key == 'post_count':
            result = by_post_count.rename({'tags': 'tag'})
        elif sort_key == 'note_count':
            result = grouped_tags.agg(pl.sum('note_count'))
            result = result.join(other=by_post_count, on='tags', how='inner')
            result = result.select(
                (pl.col('tags').alias('tag')),
                (pl.col('note_count')),
                (pl.col('post_count')),
                (pl.col('note_count').truediv(pl.col('post_count'))).alias(
                    "notes_to_posts_ratio")
            )
            result = result.sort('note_count', descending=True)
        else:
            print(f"WARN: {sort_key} is not a supported sort key!")
        return result.to_dicts()

    def get_original_posts(self) -> pl.DataFrame:
        if self.posts.is_empty():
            return pl.DataFrame()
        return self.posts.filter(pl.col('parent_post_url').is_null())

    def calculate_total_posts(self) -> int:
        if self.posts.is_empty():
            return 0
        return len(self.posts)

    def calculate_total_original_posts(self) -> int:
        if self.posts.is_empty():
            return 0
        return len(self.get_original_posts())

    def calculate_original_post_notes(self) -> int:
        if self.posts.is_empty():
            return 0
        og_posts = self.get_original_posts()
        return int(og_posts['note_count'].sum())

    def rank_original_posts(self) -> List[Dict[str, Any]]:
        og_posts = self.get_original_posts()

        sorted_posts = og_posts.clone()
        if not sorted_posts.is_empty():
            sorted_posts = sorted_posts.sort(by='note_count', descending=True)
        return sorted_posts.to_dicts()

    def determine_top_interactors(self) -> List[Dict[str, Any]]:
        if self.activity.is_empty(): return []
        og_activity = self.activity.filter(pl.col('post').struct.field('is_og') == True)
        interactor_group_by = og_activity.group_by(pl.col('from_tumblelog_name'))
        agg = interactor_group_by.agg(pl.count().alias('interaction_count')).sort('interaction_count', descending=True)
        return agg.to_dicts()

    def determine_ordered_queue(self) -> List[Dict[str, Any]]:
        if self.queued_posts.is_empty():
            return []
        sorted_queued_posts = self.queued_posts.sort(
            by='scheduled_publish_time', descending=False)
        return sorted_queued_posts.to_dicts()

    def collect_note_stats_by_original_post_date(self, date_agg: DateAgg) -> List[Dict[str, Any]]:
        og_posts = self.get_original_posts()
        if og_posts.is_empty(): return []

        sorted_og_posts = og_posts.sort(by='date')

        group_key = self._get_group_key(sorted_og_posts, date_agg)
        grouped_posts = sorted_og_posts.group_by(group_key)
        base_stats = grouped_posts.agg(pl.col('id_string').n_unique().alias('post_count'),
                                pl.sum('note_count').alias('note_count'))

        # For breakdown by tag.
        exploded_tags = sorted_og_posts.explode(
            'tags', keep_nulls=False, empty_as_null=False)
        exploded_tags = exploded_tags.rename({'tags': 'tag'})
        grouped_tag_explosion = exploded_tags.group_by([group_key, pl.col('tag')])
        by_tags = grouped_tag_explosion.agg(
            pl.col('id_string').n_unique().alias('post_count'),
            pl.sum('note_count'),
            pl.sum('note_count').truediv(pl.col('id_string').n_unique())
            .alias("notes_to_posts_ratio")).sort('note_count', descending=True)
        by_tags = by_tags.group_by(pl.col('group_key')).agg(
            pl.struct(pl.all()).alias('tags'))

        # For breakdown by post.
        grouped_by_posts = sorted_og_posts.group_by([group_key,
                                                        pl.col('post_url').alias('url')])
        by_posts = grouped_by_posts.agg(pl.sum('note_count')) \
                                        .sort('note_count', descending=True)
        by_posts = by_posts.group_by(pl.col('group_key')).agg(pl.struct(pl.all()).alias('posts'))

        post_cols = [(pl.col('note_count')),
                    (pl.col('post_count')),
                    (pl.col('note_count').truediv(pl.col('post_count'))).alias(
                        "notes_to_posts_ratio"),
                    (pl.col('tags')),
                    (pl.col('posts'))]

        # Join to get final result.
        result = base_stats.join(other=by_tags, on='group_key', how='inner')
        result = result.join(other=by_posts, on='group_key', how='inner')
        return self._date_agg_final_select(result, post_cols, date_agg).to_dicts()

    def collect_note_stats_by_original_post_year_month_day(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_original_post_date(DateAgg.DAY)

    def collect_note_stats_by_original_post_month_and_year(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_original_post_date(DateAgg.MONTH)

    def collect_note_stats_by_original_post_qtr_and_year(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_original_post_date(DateAgg.QUARTER)

    def collect_note_stats_by_original_post_year(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_original_post_date(DateAgg.YEAR)

    def collect_note_stats_by_op_interaction_month_and_year(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_op_interaction_date(DateAgg.MONTH)

    def collect_note_stats_by_op_interaction_year_month_day(self) -> List[Dict[str, Any]]:
        return self.collect_note_stats_by_op_interaction_date(DateAgg.DAY)

    def collect_note_stats_by_op_interaction_date(self, date_agg: DateAgg) -> List[Dict[str, Any]]:
        if self.activity.is_empty(): return []

        og_activity = self.activity.filter(pl.col('post').struct.field('is_og') == True)
        og_activity = og_activity.with_columns(pl.col('post').struct.field('tags').alias('og_post_tags'),
                                               pl.col('post').struct.field('url').alias('og_post_url'))
        sorted_og_activity = og_activity.sort(by='date')

        group_key = self._get_group_key(sorted_og_activity, date_agg)
        grouped_activity = sorted_og_activity.group_by(group_key)
        base_stats = grouped_activity.agg(pl.col('id').n_unique().alias('interaction_count'))
        base_stats.write_json('lesbo.json')

        # For breakdown by tag.
        exploded_tags = sorted_og_activity.explode(columns='og_post_tags', keep_nulls=False, empty_as_null=False)
        exploded_tags = exploded_tags.rename({'og_post_tags': 'tag'})
        grouped_tag_explosion = exploded_tags.group_by([group_key, pl.col('tag')])
        by_tags = grouped_tag_explosion.agg(pl.col('id').n_unique().alias('interaction_count')) \
                                        .sort('interaction_count', descending=True)
        by_tags = by_tags.group_by(pl.col('group_key')).agg(pl.struct(pl.all()).alias('tags'))


        # For breakdown by post.
        grouped_by_posts = sorted_og_activity.group_by([group_key,
                                                        pl.col('og_post_url').alias('url')])
        by_posts = grouped_by_posts.agg(pl.col('id').n_unique().alias('interaction_count')) \
                                        .sort('interaction_count', descending=True)
        by_posts = by_posts.group_by(pl.col('group_key')).agg(pl.struct(pl.all()).alias('posts'))

        actv_columns = [(pl.col('interaction_count')),
                    (pl.col('tags')),
                    (pl.col('posts'))]

        # Join to get final result.
        result = base_stats.join(other=by_tags, on='group_key', how='inner')
        result = result.join(other=by_posts, on='group_key', how='inner')
        return self._date_agg_final_select(result, actv_columns, date_agg).to_dicts()

    def _get_group_key(self, df: pl.DataFrame, date_agg: DateAgg):
        base_col = pl.col('date')
        if df.get_column('date').dtype == pl.String:
            base_col = pl.when(base_col.str.contains('GMT')) \
                .then(base_col.str.to_date('%Y-%m-%d %H:%M:%S GMT')) \
                .otherwise(base_col.str.to_date('%Y-%m-%d %H:%M:%S'))

        match date_agg:
            case DateAgg.DAY:
                group_key = base_col.dt.date().alias('group_key')
            case DateAgg.MONTH:
                group_key = base_col.dt.month_start().dt.date().alias('group_key')
            case DateAgg.QUARTER:
                group_key = pl.struct(base_col.dt.year().alias('year'),
                                                                   base_col.dt.quarter().alias('quarter')) \
                                                                   .alias('group_key')
            case DateAgg.YEAR:
                group_key = base_col.dt.year().alias('group_key')
        return group_key

    def _date_agg_final_select(self, result: pl.DataFrame, cols: List[pl.Expr], date_agg: DateAgg) -> pl.DataFrame:
        match date_agg:
            case DateAgg.DAY:
                final_cols = [(pl.col('group_key').dt.year().alias('year')),
                                    (pl.col('group_key').dt.month().alias('month')),
                                    (pl.col('group_key').dt.day().alias('day'))]
                final_cols.extend(cols)
                return result.select(final_cols).sort(by=['year', 'month', 'day'], descending=True)
            case DateAgg.MONTH:
                final_cols = [(pl.col('group_key').dt.year().alias('year')), (pl.col('group_key').dt.month().alias('month'))]
                final_cols.extend(cols)
                return result.select(final_cols).sort(by=['year', 'month'], descending=True)
            case DateAgg.QUARTER:
                final_cols = [(pl.col('group_key').alias('quarter'))]
                final_cols.extend(cols)
                return result.select(final_cols).sort(by='quarter', descending=True)
            case DateAgg.YEAR:
                final_cols = [(pl.col('group_key').alias('year'))]
                final_cols.extend(cols)
                return result.select(final_cols).sort(by='year', descending=True)