#!/usr/bin/env python3

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

TVH = "http://127.0.0.1:9981"

M3U = Path("/srv/iptv/channels-organized.m3u")

API_BASE = "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","OM","PK","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(filename):
    print(f"Downloading {filename}...")
    with urllib.request.urlopen(
            f"{API_BASE}/{filename}",
        timeout=60
    ) as response:
        return json.load(response)


def tvh_get(path, params=None):
    url = TVH + path

    if params:
        url += "?" + urllib.parse.urlencode(params)

    with urllib.request.urlopen(url) as response:
        return json.load(response)


def tvh_post(path, params):
    data = urllib.parse.urlencode(params).encode()

    request = urllib.request.Request(
        TVH + path,
        data=data
    )

    with urllib.request.urlopen(request) as response:
        return json.load(response)


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
}


def parse_attr(line, attr):
    match = re.search(
            rf'{re.escape(attr)}="([^"]*)"',
        line
    )

    return match.group(1) if match else ""


def clean_tvg_id(value):
    # iptv-org playlist IDs can include feed/quality suffixes.
    # Example:
    # FoxSports1.us@HD -> FoxSports1.us
    return value.split("@", 1)[0]


print("Parsing M3U...")

m3u_lookup = {}

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

for line in lines:
    if not line.startswith("#EXTINF"):
        continue

    tvg_id = clean_tvg_id(
        parse_attr(line, "tvg-id")
    )

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

    if not name:
        continue

    meta = channel_meta.get(tvg_id)

    if meta:
        country_code = meta.get("country") or ""
        categories = meta.get("categories") or []

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

        elif categories:
            genre = category_names.get(
                categories[0],
                categories[0].replace("_", " ").title()
            )

        else:
            genre = "Undefined"

    else:
        country_code = ""
        genre = "Undefined"

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

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

    entry = {
        "tvg_id": tvg_id,
        "continent": continent,
        "country": country,
        "genre": genre,
        "name": name,
    }

    # Multiple streams can legitimately have the same visible name.
    m3u_lookup.setdefault(
        name.casefold(),
        []
    ).append(entry)


print("Loading Tvheadend channels...")

tvh_channels = tvh_get(
    "/api/channel/grid",
    {
        "limit": 50000,
        "all": 1,
    }
)["entries"]


continent_rank = {
    value: index
    for index, value in enumerate(CONTINENT_ORDER)
}


resolved = []
unresolved = []

for channel in tvh_channels:

    name = channel.get("name", "").strip()

    candidates = m3u_lookup.get(
        name.casefold(),
        []
    )

    if len(candidates) == 1:

        meta = candidates[0]

    elif len(candidates) > 1:

        # Same-name streams exist in the world list.
        # For sorting purposes, take the deterministic
        # first candidate by location/genre/id.
        meta = sorted(
            candidates,
            key=lambda x: (
                continent_rank.get(
                    x["continent"], 999
                ),
                x["country"].casefold(),
                x["genre"].casefold(),
                x["tvg_id"].casefold(),
            )
        )[0]

    else:

        unresolved.append(name)

        meta = {
            "continent": "Other / International",
            "country": "Unknown",
            "genre": "Undefined",
            "name": name,
            "tvg_id": "",
        }

    resolved.append({
        "uuid": channel["uuid"],
        "name": name,
        **meta,
    })


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


# Sequential numbering is enough.
#
# The number itself doesn't encode the hierarchy;
# it merely preserves the already-correct sort order.
changes = []

for number, channel in enumerate(
        resolved,
    start=1
):
    changes.append({
        **channel,
        "number": number,
    })


print()
print(f"Tvheadend channels: {len(changes):,}")
print(f"Unresolved names:    {len(unresolved):,}")

print()
print("First 60 proposed assignments:")
print()

for x in changes[:60]:
    print(
        f'{x["number"]:6d}  '
        f'{x["continent"]} | '
        f'{x["country"]} | '
        f'{x["genre"]} | '
        f'{x["name"]}'
    )


if unresolved:

    print()
    print("First 20 unresolved channels:")

    for name in unresolved[:20]:
        print(f"  {name}")


answer = input(
    "\nApply these channel numbers? [y/N] "
).strip().lower()

if answer != "y":
    print("No changes made.")
    raise SystemExit(0)


print()
print("Updating Tvheadend...")

for index, channel in enumerate(
        changes,
    start=1
):

    node = json.dumps({
        "uuid": channel["uuid"],
        "number": channel["number"],
    })

    tvh_post(
        "/api/idnode/save",
        {
            "node": node
        }
    )

    if index % 250 == 0:
        print(
            f"Updated "
            f"{index:,} / "
            f"{len(changes):,}"
        )


print()
print(
    f"Done. Numbered "
    f"{len(changes):,} channels."
)