import pytz
from datetime import datetime
import pandas as pd


def get_tz_info(tz_name: str, timezones: list[dict]) -> dict | None:
    """Get timezone info by name."""
    return next((tz for tz in timezones if tz["zone_name"] == tz_name), None)


def get_country_info(country_code: str, countries: list[dict]) -> dict | None:
    """Get country info by country code."""
    return next((c for c in countries if c["country_code"] == country_code), None)


def where_is_it_420(timezones: list[dict], countries: list[dict]) -> list[str]:
    """Get timezones where the current hour is 4 or 16, indicating it's 4:20 there.

    Returns:
        list[str]: A list of timezones where it's currently 4:20 PM or AM.
    """
    tz_list = []
    for tz in pytz.all_timezones:
        now = datetime.now(pytz.timezone(tz))
        if now.hour == 4 or now.hour == 16:
            # Find the timezone in the loaded timezones
            tz_info = get_tz_info(tz, timezones)
            if tz_info:
                country = get_country_info(tz_info["country_code"], countries)
                if country:
                    country_name = country["country_name"].strip().strip('"')
                    if country_name not in tz_list:
                        tz_list.append(country_name)
    return tz_list


def load_tz_file():
    timezone_file = "./tzdb/TimeZoneDB.csv/time_zone.csv"
    # column names in the csv
    timezone_names = ["zone_name", "country_code",
                      "abbreviation", "time_start", "gmt_offset", "dst"]
    # columns to load
    load_columns = ["zone_name", "country_code"]
    # read csv with pandas
    df = pd.read_csv(timezone_file, names=timezone_names)
    # drop all columns except load_columns
    df = df[load_columns]
    # distinct zone_names
    df = df.drop_duplicates(subset=["zone_name"])

    # reset index
    df = df.reset_index(drop=True)

    return df


def main():

    # read csv with pandas
    df_file = load_tz_file()

    # split zone_name into components by "/"
    df_file[['region', 'city']] = df_file['zone_name'].str.split(
        '/', expand=True, n=1)
    # drop regions with no country_code (like Etc, GMT, etc)
    df_file = df_file[df_file['country_code'].notna()]

    df_tz = pd.DataFrame(pytz.all_timezones)
    # rename column to zone_name
    df_tz = df_tz.rename(columns={0: 'zone_name'})
    # split zone_name into components by "/"
    df_tz[['region', 'city']] = df_tz['zone_name'].str.split(
        '/', expand=True, n=1)
    # drop regions with no city (like UTC, GMT, etc)
    df_tz = df_tz[df_tz['city'].notna()]
    # drop rows where region is 'Etc'
    df_tz = df_tz[df_tz['region'] != 'Etc']

    # join dataframes on region and city
    df_merged = pd.merge(df_file, df_tz, on=[
                         'region', 'city'], how='inner', indicator=True)
    # reorder columns
    df_merged = df_merged[['region', 'city', 'country_code']]
    # print merged dataframe
    print(f"Merged timezones: {len(df_merged)}")
    print(df_merged.sample(20).to_string(index=False))
    regions = df_merged['region'].unique()
    for region in regions:
        df_region = df_merged[df_merged['region'] == region]
        print(f"{len(df_region)} merged in {region}")


if __name__ == "__main__":
    main()