import json
import logging
import os
import time
from datetime import datetime, timedelta

from apollo_main_db.push_notifications.models import PushMessageReduced
from constants import BATCH_SIZE, CHART_TYPES, PLAYLIST_TYPES, TOPIC_MAP
from mysql_config import session as mysql_session
from psql_config import session_scope as psql_session
from sqlalchemy import and_
from user_data_tables import Account, FeedMessage


def get_dates():
    start_date = os.environ.get('START_DATETIME')
    end_date = os.environ.get('END_DATETIME')
    if start_date:
        start_date = datetime.strptime(start_date, '%Y-%m-%d')
    else:
        start_date = datetime.utcnow() - timedelta(days=30)
    if end_date:
        end_date = datetime.strptime(end_date, '%Y-%m-%d')
    else:
        end_date = datetime.utcnow() - timedelta(days=1)
    return start_date, end_date


def get_query():
    start_date, end_date = get_dates()
    query = mysql_session.query(
        PushMessageReduced.id,
        PushMessageReduced.inner_id,
        PushMessageReduced.user_id,
        PushMessageReduced.title,
        PushMessageReduced.message,
        PushMessageReduced.data,
        PushMessageReduced.track_id,
        PushMessageReduced.topic,
        PushMessageReduced.created_at,
        PushMessageReduced.is_new,
        PushMessageReduced.vendor,
        PushMessageReduced.date,
    ).filter(
        and_(PushMessageReduced.created_at >= start_date, PushMessageReduced.created_at < end_date)
    ).order_by(PushMessageReduced.id)
    return query


def backfill_notifications():
    query = get_query()

    loaded_total = 0

    offset = 0
    while True:
        records = query.slice(offset, offset + BATCH_SIZE).all()

        start_time = time.time()

        logging.info(f"Processing batch of {len(records)}")

        ids = [i["inner_id"] for i in records]

        with psql_session() as session:
            existing_records = session.query(FeedMessage.message_id).filter(FeedMessage.message_id.in_(ids)).all()
            existing_records = [i[0] for i in existing_records]
            for item in [i for i in records if i["inner_id"] not in existing_records]:
                item = item._asdict()
                inner_id = item["inner_id"]
                item_data = json.loads(item["data"])

                meta = {
                    "domain": "feed_message",
                    "dsp": item["vendor"].name,
                    "country_code": item_data.get("country_code"), 
                    "type": TOPIC_MAP.get(item["topic"].name)
                }

                if item["topic"].name in CHART_TYPES:
                    if item["topic"].name == "CHART_ADDITIONS":
                        current_position = item_data.get("position")
                        previous_position = None
                    elif item["topic"].name == "CHART_REMOVALS":
                        current_position = None
                        previous_position = item_data.get("position")
                    else:
                        current_position = item_data.get("position")
                        previous_position = item_data.get("position") - item_data.get("change")
                    data = {
                        "content": {
                            "track": {
                                "id": item_data.get("track_id"),
                                "isrc": item_data.get("isrc"),
                                "name": item_data.get("track_name"),
                                "artists": [{"id": None, "name": item_data.get("artist_name")}],
                            },
                            "chart": {
                                "name": item_data.get("target"),
                            },
                            "body": item["message"],
                            "title": item["title"],
                            "current_position": current_position,
                            "previous_position": previous_position,
                        }
                    }

                elif item["topic"].name in PLAYLIST_TYPES:
                    data = {
                        "content": {
                            "track": {
                                "id": item_data.get("track_id"),
                                "isrc": item_data.get("isrc"),
                                "name": item_data.get("track_name"),
                                "artists": [{"id": None, "name": item_data.get("artist_name")}],
                            },
                            "playlist": {
                                "id": item_data.get("playlist_id"),
                                "name": item_data.get("playlist_name"),
                                "image_url": item_data.get("playlist_image_url"),
                            },
                            "body": item["message"],
                            "title": item["title"],
                            "current_position": item_data.get("position"),
                            "previous_position": None,
                        }
                    }
                
                else:
                    data = {
                        "content": {
                            "playlist": {
                                "id": item_data.get("playlist_id"),
                                "name": item_data.get("playlist_name"),
                                "image_url": item_data.get("playlist_image_url"),
                                "updated_at": item_data.get("date"),
                            },
                            "body": item["message"],
                            "title": item["title"],
                        },
                    }
            
                account_id = session.query(Account.id).filter(
                    Account.app_slug == "apollo",
                    Account.user_id == item["user_id"]
                ).first()[0]
                session.add(FeedMessage(
                    created_at=item["created_at"],
                    updated_at=datetime.utcnow(),
                    public=False,
                    ttl=30 * 60,
                    message_id=inner_id,
                    event_id=None,
                    status=item["is_new"],
                    app_slug="apollo",
                    meta=meta,
                    data=data,
                    account_id=account_id
                ))
        session.commit()

        loaded_total += len(records)
            
        logging.info(f"Batch completed within {(time.time() - start_time)}")
        logging.info(f"Skipped {len(existing_records)} existing records")
        logging.info(f"{loaded_total} in total records processed")

        if len(records) < BATCH_SIZE:
            break
        else:
            offset += BATCH_SIZE


if __name__ == '__main__':
    backfill_notifications()