import argparse
import cloudscraper
import logging
from bs4 import BeautifulSoup
import pandas as pd
import os
import threading
from datetime import datetime, timedelta

from slz_spotify_charts_scraper.const import SPOTIFY_TRACK_PREFIX


REGIONAL_CHART = 'regional'
VIRAL_CHART = 'viral'
CHART_TYPES = (REGIONAL_CHART, VIRAL_CHART)

DAILY_CHART = 'daily'
WEEKLY_CHART = 'weekly'
RECURRENCE_TYPES = (DAILY_CHART, WEEKLY_CHART)

logging.basicConfig(level=logging.DEBUG)


def get_market_codes(chart_type, recurrence_type, debug=False):
    url = f'https://spotifycharts.com/{chart_type}/global/{recurrence_type}/latest'
    scraper = cloudscraper.create_scraper(debug=debug)
    
    page = scraper.get(url, stream=True)
    soup = BeautifulSoup(page.content, "html.parser")
    markets_selector_container = soup.find("div", attrs={"data-type": "country"})
    markets_list_items = markets_selector_container.find_all("li")
    return [item["data-value"] for item in markets_list_items if item and item.get("data-value")] if markets_list_items else []

def get_chart_tracks(chart_type, recurrence_type, market, date_str, debug=False):
    url = f'https://spotifycharts.com/{chart_type}/{market}/{recurrence_type}/{date_str}'
    scraper = cloudscraper.create_scraper(debug=debug)    
    page = scraper.get(url, stream=True)

    if page.status_code >= 500:
        logging.error('Server Error!')

    soup = BeautifulSoup(page.content, 'html.parser')

    try:
        chart_table = soup.find('table', {'class': 'chart-table'})
        songs = chart_table.find("tbody").find_all("tr")

        result = []
        for song in songs:
            image_container = song.find("td", {"class": 'chart-table-image'})
            uri = image_container.find("a")["href"]
            data = {
                'position': int(song.find("td", {"class": "chart-table-position"}).text),
                'name': song.find("td", {"class": "chart-table-track"}).find("strong").text,
                'image': song.find("td", {"class": "chart-table-image"}).find("img")["src"],
                'artist': song.find("td", {"class": "chart-table-track"}).find("span").text[3:],
                'uri': uri,
                'id': uri.replace(SPOTIFY_TRACK_PREFIX, '')
            }
            if chart_type == REGIONAL_CHART:
                data["streams"] = int(song.find("td", {"class": "chart-table-streams"}).text.replace(",", ""))
            result.append(data)
        return pd.DataFrame(result)

    except Exception:
        return pd.DataFrame()

def get_date_string_parameter(target_date, chart_type, recurrence_type):
    if recurrence_type == DAILY_CHART:
        return target_date.isoformat()

    if chart_type == VIRAL_CHART:
        date_str = target_date.isoformat()
        return f"{date_str}--{date_str}"

    end_date = target_date + timedelta(days=1)
    start_date = target_date - timedelta(days=6)
    return f"{start_date.isoformat()}--{end_date.isoformat()}"

def save_charts_file(chart_type, recurrence_type, market, date):
    date_param = get_date_string_parameter(date, chart_type, recurrence_type)
    df = get_chart_tracks(chart_type, recurrence_type, market, date_param)
    if df.empty:
        return
    filename = f"./results/spotify/charts/{chart_type}/{recurrence_type}/{date}/charts_{market}.parquet"
    os.makedirs(os.path.dirname(filename), exist_ok=True)
    df.to_parquet(filename, engine="pyarrow")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--chart",
        type=str,
        required=True,
        choices=CHART_TYPES,
    )
    parser.add_argument(
        "-r",
        "--recurrence",
        type=str,
        required=True,
        choices=RECURRENCE_TYPES,
    )
    parser.add_argument('--debug', action='store_true')
    parser.add_argument(
        '-td',
        '--target-date',
        type=str,
        required=True,
        help='start report date %YYYY-%MM-%DD (e.g. 2021-03-01)'
    )
    args = parser.parse_args()
    markets = get_market_codes(args.chart, args.recurrence)
    target_date = datetime.strptime(args.target_date, "%Y-%m-%d").date()

    threads = []
    for market in markets:
        th = threading.Thread(target=save_charts_file, args=(args.chart, args.recurrence, market, target_date))
        threads.append(th)
        th.start()

    for thread in threads:
        thread.join()


if __name__ == "__main__":
    main()