from datetime import datetime as dt
import json
import sys
from time import sleep
import traceback
from typing import Any, Dict, List

import pytumblr
import polars as pl


def filter_for_after(df: pl.DataFrame,
                           after: dt) -> pl.DataFrame:
    if df.is_empty() or not after: return pl.DataFrame()
    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'))
    return df.filter(base_col >= after.date())


def filter_for_before(df: pl.DataFrame,
                            before: dt) -> pl.DataFrame:
    if df.is_empty() or not before: return pl.DataFrame()
    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'))
    return df.filter(base_col <= before.date())


def handle_error(e: Exception, data: Any, df: pl.DataFrame) -> None:
    print(e)
    traceback.print_exc()
    schema_dict = {col: str(dtype) for col, dtype in df.schema.items()}
    with open("schema.json", "w") as f:
        json.dump(schema_dict, f, indent=4)
    with open("bad_posts.json", "w") as f:
        json.dump(data, f, indent=4, default=str)


def handle_limit_exceeded(data) -> bool:
    if 'meta' in data and data['meta']['status'] == 429:
        print("WARN: API limit exceeded. Sleeping for 10 minutes...")
        sleep(600)
        return True
    sleep(3)
    return False


def fetch_minimized_post(blog_name: str,
                           client: pytumblr.TumblrRestClient,
                           id_string: str) -> Dict[str, Any]:
    api_url = f"/v2/blog/{blog_name}/posts"
    data: Dict[str, Any] = client.send_api_request(method='get',
                                   url=api_url,
                                   params={'id': id_string},
                                   valid_parameters=['id'])

    if 'posts' in data:
        posts = data['posts']
    if handle_limit_exceeded(data):
        data = client.send_api_request(method='get', url=api_url, params={'id': id_string}, valid_parameters=['id'])
        posts = data['posts']
    if 'posts' not in data or not data['posts']:
        print(f"WARN: Post with ID {id_string} could not be retrieved!")
        print(data)
        return {}

    minimized_post: Dict[str, Any] = {}
    if posts:
        minimized_post = posts[0]
        minimized_post = {
            'id_string': minimized_post['id_string'],
            'url': minimized_post['post_url'],
            'date_str': minimized_post['date'],
            'tags': minimized_post['tags'],
            'is_og': 'parent_post_url' not in minimized_post
        }
    return minimized_post


def fetch_queued_posts(client: pytumblr.TumblrRestClient,
                     args: Dict[str, Any]) -> pl.DataFrame:
    queued_posts: pl.DataFrame = pl.DataFrame()
    blog_name: str = args['blog']

    offset: int = 0
    limit: int = 20

    while True:
        # Begin LOOP
        # Get me some queued posts via REST! 😈🍪🍪🍪
        data: Dict[str, Any]
        data = client.queue(f"{blog_name}.tumblr.com",
                            offset=offset,
                            limit=limit)

        if handle_limit_exceeded(data):
            data = client.queue(f"{blog_name}.tumblr.com",
                    offset=offset,
                    limit=limit)

        # This struct breaks Polars for some reason 🤷
        posts_json = data['posts']
        for post in posts_json:
            if 'trail' in post:
                del post['trail']

        try:
            curr_posts: pl.DataFrame = pl.DataFrame(posts_json)
        except Exception as e:
            handle_error(e, data, queued_posts)
            sys.exit(0)

        # Stop the presses if we found no posts.
        if curr_posts.is_empty() or len(curr_posts) < 1:
            print('Stopping queued post retrieval, as no more queued posts were found.')
            break

        queued_posts = pl.concat([queued_posts, curr_posts], how="diagonal_relaxed")

        offset += limit
        # End LOOP

    # Empty check.
    if queued_posts.is_empty():
        print("No queued posts were found.")

    return queued_posts


def fetch_drafts(client: pytumblr.TumblrRestClient,
               args: Dict[str, Any]) -> pl.DataFrame:
    drafts: pl.DataFrame = pl.DataFrame()
    blog_name: str = args['blog']
    draft_url = f"/v2/blog/{blog_name}/posts/draft"

    while True:
        # Begin LOOP
        # Get me some drafts via REST! 😈🍪🍪🍪
        data: Dict[str, Any]
        data = client.send_api_request("get", draft_url)

        if handle_limit_exceeded(data):
            data = client.send_api_request("get", draft_url)

        # This struct breaks Polars for some reason 🤷
        drafts_json = data['posts']
        for draft in drafts_json:
            if 'trail' in draft:
                del draft['trail']

        curr_drafts: pl.DataFrame = pl.DataFrame(drafts_json)
        drafts = pl.concat([drafts, curr_drafts], how="diagonal_relaxed")

        if '_links' in data:
            draft_url = data['_links']['next']['href']
            continue
        else:
            print('Stopping draft retrieval, as no more drafts were found.')
            break
        # End LOOP

    # Empty check.
    if drafts.is_empty():
        print("No drafts were found.")

    return drafts


def fetch_posts(client: pytumblr.TumblrRestClient,
              args: Dict[str, Any]) -> pl.DataFrame:
    posts: pl.DataFrame = pl.DataFrame()
    blog_name: str = args['blog']

    # We populate params, starting with any tags for filtering.
    params = {}
    if args['tags']:
        params.update({'tag': ','.join(args['tags'])})

    total: int = 0
    offset: int = 0
    limit: int = 20

    # The request loop that pulls all data from the APIs.
    while True:
        # Begin LOOP
        # Get me some posts via REST! 😈🍪🍪🍪
        data: Dict[str, Any]
        data = client.posts(f"{blog_name}.tumblr.com",
                            offset=offset,
                            limit=limit,
                            **params)

        if handle_limit_exceeded(data):
            data = client.posts(f"{blog_name}.tumblr.com",
                    offset=offset,
                    limit=limit,
                    **params)

        # These structs break Polars for some reason 🤷
        posts_json = data['posts']
        for post in posts_json:
            if 'trail' in post:
                del post['trail']
            if 'player' in post:
                del post['player']

        try:
            curr_posts: pl.DataFrame = pl.DataFrame(posts_json, strict=False, infer_schema_length=None)
        except Exception as e:
            handle_error(e, data, curr_posts)
            sys.exit(1)

        # Stop the presses if we found no posts.
        if curr_posts.is_empty() or len(curr_posts) < 1:
            print('Stopping, as no more posts were found.')
            break
        
        next_off: int = 0

        # Total init check for the first iteration, but always checked for sanity.
        if not total and 'total_posts' in data:
            total_posts = data['total_posts']
            print(f"I'm working with {total_posts} total posts...")
            total = total_posts

        if args['since'] and not args['until']:
            after: dt = args['since']
            curr_posts = filter_for_after(curr_posts, after)
            if curr_posts.is_empty():
                print(
                    f"All posts since {after.year}-{after.month}-{after.day} retrieved.")
                return posts
        if args['until'] and not args['since']:
            before: dt = args['until']
            curr_posts = filter_for_before(curr_posts, before)
            if curr_posts.is_empty():
                print(
                    f"All posts until {before.year}-{before.month}-{before.day} retrieved.")
                return posts
        if args['since'] and args['until']:
            if not posts.is_empty() and curr_posts.is_empty():
                print(f"All posts since {after.year}-{after.month}-{after.day}, "
                      + f"until {before.year}-{before.month}-{before.day} retrieved.")
                return posts

        # Update the frame with what we found.
        try:
            posts = pl.concat([posts, curr_posts], how="diagonal_relaxed")
        except Exception as e:
            handle_error(e, data, posts)
            sys.exit(1)

        # The increment and status printing.
        if next_off != 0 and next_off != offset:
            offset = next_off
        else:
            offset += limit
        if not (args['since'] or args['until']) and total:
            print(
                f"Processed batch {offset // limit} of {(total // 20) + 1}...")
        # End LOOP

    if posts.is_empty():
        print("No posts were found.")

    # Return posts (both original and not original).
    return posts.unique(subset=["id_string"])


def fetch_activity(client: pytumblr.TumblrRestClient,
              args: Dict[str, Any],
              posts: pl.DataFrame = pl.DataFrame()) -> pl.DataFrame:
    activity: pl.DataFrame = pl.DataFrame(infer_schema_length=None)
    blog_name: str = args['blog']

    before: dt = args['until']
    after: dt = args['since']

    activity_url = f"/v2/blog/{blog_name}/notifications"
    if before:
        activity_url += f"?before={str(int(before.timestamp()))}"

    params = {}
    params.update({'rollups': 'false'})

    post_map: Dict[str, Any] = {}
    if not posts.is_empty():
         post_map = { row["id"]: row for row in posts.to_dicts() }

    data: Dict[str, Any]
    while True:
        # Begin LOOP
        # Get me some activity via REST! 😈🍪🍪🍪
        data = client.send_api_request("get", activity_url)
        if handle_limit_exceeded(data):
             data = client.send_api_request("get", activity_url)

        try:
            curr_activity: pl.DataFrame = pl.DataFrame(data['notifications'], infer_schema_length=None)
        except Exception as e:
            handle_error(e, data, activity)
            sys.exit(0)
        curr_activity = curr_activity.with_columns(pl.from_epoch('timestamp').alias('date'))
        if before:
            curr_activity = filter_for_before(curr_activity, before)
        if after:
            curr_activity = filter_for_after(curr_activity, after)

        # Add additional post info to notifications.
        # Yeah, df to dictionary list. Sorry. :c
        curr_activity_dicts: List[Dict[str, Any]] = curr_activity.to_dicts()
        for notification in curr_activity_dicts:
            if 'target_post_id' in notification and notification['target_post_id']:
                target_post_id = notification['target_post_id']               
                post: Any
                if target_post_id in post_map:
                    post = post_map.get(target_post_id)
                    notification['post'] = post
                    continue
                elif 'type' in notification and notification['type'] == 'post_attribution':
                    notification['post'] = {}
                    continue
                else:
                    post = fetch_minimized_post(blog_name, client, notification['target_post_id'])
                    notification['post'] = post
                    post_map.update({target_post_id: post})
                    continue

        curr_activity = pl.DataFrame(curr_activity_dicts, strict=False, infer_schema_length=None)
        activity = pl.concat([activity, curr_activity], how="diagonal_relaxed")

        if '_links' in data:
            next_time =  dt.fromtimestamp(int(data['_links']['next']['query_params']['before']))
            if before or after:
                print(f"Reached activity time: {next_time}...")
                if (before and next_time >= before) or (after and next_time <= after):
                    print('Stopping activity retrieval, as no more activity is left in date range.')
                    break
            activity_url = data['_links']['next']['href']
            continue
        else:
            print('Stopping activity retrieval, as no more activity was found.')
            break
        # End LOOP

    # Empty check.
    if activity.is_empty():
        print("No activity was found.")
    else:
        activity = activity.unique(subset=["id"])

    return activity