#!/usr/bin/env python3

import json
import re
import urllib.request
from pathlib import Path

SRC = Path("/srv/iptv/channels-organized.m3u")
OUT = Path("/srv/iptv/channels-sorted.m3u")

API = "https://iptv-org.github.io/api"

CONTINENT_ORDER = [
    "North America",
    "South America",
    "Europe",
    "Africa",
    "Asia",
    "Oceania",
    "Other / International",
]

COUNTRY_TO_CONTINENT = {}

for code in [
    "US","CA","MX","GL","BM","BZ","CR","SV","GT","HN","NI","PA",
    "AI","AG","AW","BS","BB","BQ","VG","KY","CU","CW","DM","DO",
    "GD","GP","HT","JM","MQ","MS","PR","BL","KN","LC","MF","VC",
    "SX","TT","TC","VI"
]:
    COUNTRY_TO_CONTINENT[code] = "North America"

for code in [
    "AR","BO","BR","CL","CO","EC","FK","GF","GY","PY","PE","SR","UY","VE"
]:
    COUNTRY_TO_CONTINENT[code] = "South America"

for code in [
    "AL","AD","AT","BY","BE","BA","BG","HR","CY","CZ","DK","EE","FO",
    "FI","FR","DE","GI","GR","GG","HU","IS","IE","IM","IT","JE","LV",
    "LI","LT","LU","MT","MD","MC","ME","NL","MK","NO","PL","PT","RO",
    "RU","SM","RS","SK","SI","ES","SE","CH","UA","UK","VA"
]:
    COUNTRY_TO_CONTINENT[code] = "Europe"

for code in [
    "DZ","AO","BJ","BW","BF","BI","CV","CM","CF","TD","KM","CG","CD",
    "CI","DJ","EG","GQ","ER","SZ","ET","GA","GM","GH","GN","GW","KE",
    "LS","LR","LY","MG","MW","ML","MR","MU","YT","MA","MZ","NA","NE",
    "NG","RE","RW","SH","ST","SN","SC","SL","SO","ZA","SS","SD","TZ",
    "TG","TN","UG","EH","ZM","ZW"
]:
    COUNTRY_TO_CONTINENT[code] = "Africa"

for code in [
    "AF","AM","AZ","BH","BD","BT","BN","KH","CN","GE","HK","IN","ID",
    "IR","IQ","IL","JP","JO","KZ","KW","KG","LA","LB","MO","MY","MV",
    "MN","MM","NP","KP","PS","PH","QA","SA","SG","KR","LK","SY","TW",
    "TJ","TH","TL","TR","TM","AE","UZ","VN","YE"
]:
    COUNTRY_TO_CONTINENT[code] = "Asia"

for code in [
    "AS","AU","CK","FJ","PF","GU","KI","MH","FM","NR","NC","NZ","NU",
    "NF","MP","PW","PG","PN","WS","SB","TK","TO","TV","VU","WF"
]:
    COUNTRY_TO_CONTINENT[code] = "Oceania"


def fetch_json(name):
    print(f"Downloading {name}...")
    with urllib.request.urlopen(f"{API}/{name}", timeout=60) as r:
        return json.load(r)


def attr(line, name):
    m = re.search(rf'{re.escape(name)}="([^"]*)"', line)
    return m.group(1) if m else ""


channels_api = fetch_json("channels.json")
countries_api = fetch_json("countries.json")
categories_api = fetch_json("categories.json")

channel_meta = {x["id"]: x for x in channels_api}
country_names = {x["code"]: x["name"] for x in countries_api}
category_names = {x["id"]: x["name"] for x in categories_api}

continent_rank = {
    name: i
    for i, name in enumerate(CONTINENT_ORDER)
}

lines = SRC.read_text(
    encoding="utf-8",
    errors="replace"
).splitlines()

# Preserve everything before the first #EXTINF.
preamble = []
entries = []
current = None

for line in lines:

    if line.startswith("#EXTINF"):

        if current is not None:
            entries.append(current)

        current = {
            "extinf": line,
            "lines": [line],
        }

    elif current is None:
        preamble.append(line)

    else:
        # Preserve ALL lines belonging to this channel:
        # KODIPROP, EXTVLCOPT, comments, URL, etc.
        current["lines"].append(line)

if current is not None:
    entries.append(current)


for entry in entries:

    extinf = entry["extinf"]

    tvg_id = attr(
        extinf,
        "tvg-id"
    ).split("@", 1)[0]

    name = extinf.rsplit(",", 1)[-1].strip()

    meta = channel_meta.get(tvg_id)

    if meta:

        country_code = meta.get("country") or ""

        if meta.get("is_nsfw"):
            genre = "Adult"

        else:
            cats = meta.get("categories") or []

            genre = (
                category_names.get(
                    cats[0],
                    cats[0].replace("_", " ").title()
                )
                if cats
                else "Undefined"
            )

    else:
        country_code = ""
        genre = "Undefined"

    country = country_names.get(
        country_code,
        "Unknown"
    )

    continent = COUNTRY_TO_CONTINENT.get(
        country_code,
        "Other / International"
    )

    entry["continent"] = continent
    entry["country"] = country
    entry["genre"] = genre
    entry["name"] = name


entries.sort(
    key=lambda x: (
        continent_rank.get(
            x["continent"],
            999
        ),
        x["country"].casefold(),
        x["genre"].casefold(),
        x["name"].casefold(),
    )
)


output = list(preamble)

for entry in entries:
    output.extend(entry["lines"])


OUT.write_text(
    "\n".join(output) + "\n",
    encoding="utf-8"
)

print()
print(f"Channels written: {len(entries):,}")
print(f"Output: {OUT}")