#!/usr/bin/env python3

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

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

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

MAJOR_LANGUAGES = {
    "eng": "English",
    "spa": "Spanish",
    "fra": "French",
    "deu": "German",
    "ita": "Italian",
    "por": "Portuguese",
    "rus": "Russian",
    "ara": "Arabic",
    "hin": "Hindi",
    "ben": "Bengali",
    "urd": "Urdu",
    "tam": "Tamil",
    "tel": "Telugu",
    "mar": "Marathi",
    "pan": "Punjabi",
    "guj": "Gujarati",
    "mal": "Malayalam",
    "kan": "Kannada",
    "nep": "Nepali",
    "tur": "Turkish",
    "fas": "Persian",
    "heb": "Hebrew",
    "zho": "Chinese",
    "cmn": "Mandarin",
    "yue": "Cantonese",
    "jpn": "Japanese",
    "kor": "Korean",
    "tha": "Thai",
    "vie": "Vietnamese",
    "ind": "Indonesian",
    "msa": "Malay",
    "fil": "Filipino",
    "nld": "Dutch",
    "pol": "Polish",
    "ukr": "Ukrainian",
    "ron": "Romanian",
    "ell": "Greek",
    "swe": "Swedish",
    "nor": "Norwegian",
    "dan": "Danish",
    "fin": "Finnish",
    "ces": "Czech",
    "hun": "Hungarian",
}

CATEGORY_NAMES = {
    "animation": "Animation",
    "auto": "Auto",
    "business": "Business",
    "classic": "Classic",
    "comedy": "Comedy",
    "cooking": "Cooking",
    "culture": "Culture",
    "documentary": "Documentary",
    "education": "Education",
    "entertainment": "Entertainment",
    "family": "Family",
    "general": "General",
    "interactive": "Interactive",
    "kids": "Kids",
    "legislative": "Legislative",
    "lifestyle": "Lifestyle",
    "movies": "Movies",
    "music": "Music",
    "news": "News",
    "outdoor": "Outdoor",
    "public": "Public",
    "relax": "Relax",
    "religious": "Religious",
    "science": "Science",
    "series": "Series",
    "shop": "Shopping",
    "sports": "Sports",
    "travel": "Travel",
    "weather": "Weather",
    "xxx": "Adult",
}

# Broad Kodi-friendly regions. Country remains metadata; this is just browsing.
REGION_COUNTRIES = {
    "North America": {
        "US", "CA", "MX", "GL", "BM", "PM"
    },

    "Central America": {
        "BZ", "CR", "SV", "GT", "HN", "NI", "PA"
    },

    "Caribbean": {
        "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"
    },

    "South America": {
        "AR", "BO", "BR", "CL", "CO", "EC", "FK", "GF",
        "GY", "PY", "PE", "SR", "UY", "VE"
    },

    "Europe": {
        "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"
    },

    "Middle East": {
        "BH", "IR", "IQ", "IL", "JO", "KW", "LB", "OM",
        "PS", "QA", "SA", "SY", "TR", "AE", "YE"
    },

    "South Asia": {
        "AF", "BD", "BT", "IN", "MV", "NP", "PK", "LK"
    },

    "East Asia": {
        "CN", "HK", "MO", "JP", "KP", "KR", "MN", "TW"
    },

    "Southeast Asia": {
        "BN", "KH", "ID", "LA", "MY", "MM", "PH", "SG", "TH", "TL", "VN"
    },

    "Central Asia": {
        "KZ", "KG", "TJ", "TM", "UZ"
    },

    "Africa": {
        "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"
    },

    "Oceania": {
        "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_REGION = {}
for region, countries in REGION_COUNTRIES.items():
    for code in countries:
        COUNTRY_TO_REGION[code] = region


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


channels = get_json("channels.json")
feeds = get_json("feeds.json")
languages = get_json("languages.json")

language_names = {
    x["code"]: x["name"]
    for x in languages
}

channel_data = {
    c["id"]: c
    for c in channels
}

# Languages are feed-level in the current API.
channel_languages = defaultdict(set)

for feed in feeds:
    channel = feed.get("channel")
    if not channel:
        continue

    for lang in feed.get("languages", []):
        channel_languages[channel].add(lang)


def pretty_language(channel_id):
    langs = channel_languages.get(channel_id)

    if not langs:
        return "Unknown"

    # Prefer a recognizable/common language when multiple exist.
    for lang in sorted(langs):
        if lang in MAJOR_LANGUAGES:
            return MAJOR_LANGUAGES[lang]

    # Avoid producing 150 microscopic Kodi groups.
    return "Other"


def pretty_category(channel):
    cats = channel.get("categories") or []

    if channel.get("is_nsfw"):
        return "Adult"

    if not cats:
        return "Undefined"

    # First category is good enough for primary browsing.
    cat = cats[0]

    return CATEGORY_NAMES.get(cat, cat.replace("_", " ").title())


def get_region(country):
    return COUNTRY_TO_REGION.get(country, "Other / International")


def replace_group_title(extinf, new_group):
    if 'group-title="' in extinf:
        return re.sub(
            r'group-title="[^"]*"',
            f'group-title="{new_group}"',
            extinf,
            count=1
        )

    comma = extinf.rfind(",")

    if comma != -1:
        return (
            extinf[:comma]
            + f' group-title="{new_group}"'
            + extinf[comma:]
        )

    return extinf


def get_tvg_id(line):
    m = re.search(r'tvg-id="([^"]*)"', line)
    return m.group(1) if m else ""


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

output = []

if lines:
    output.append(lines[0])

stats = defaultdict(int)

i = 1

while i < len(lines):
    line = lines[i]

    if not line.startswith("#EXTINF"):
        output.append(line)
        i += 1
        continue

    extinf = line
    url = lines[i + 1] if i + 1 < len(lines) else ""

    tvg_id = get_tvg_id(extinf)

    # iptv-org sometimes uses IDs like Foo.us@HD
    base_id = tvg_id.split("@", 1)[0]

    channel = channel_data.get(tvg_id) or channel_data.get(base_id)

    if channel:
        country = channel.get("country") or ""
        region = get_region(country)
        language = pretty_language(channel["id"])
        category = pretty_category(channel)
    else:
        region = "Other / International"
        language = "Unknown"

        m = re.search(r'group-title="([^"]*)"', extinf)
        category = m.group(1) if m and m.group(1) else "Undefined"

    group = f"{region} | {language} | {category}"

    output.append(replace_group_title(extinf, group))
    output.append(url)

    stats[group] += 1

    i += 2


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

print()
print(f"Wrote: {OUT}")
print(f"Channels: {sum(stats.values()):,}")
print(f"Groups: {len(stats):,}")

print()
print("Largest groups:")

for group, count in sorted(
    stats.items(),
    key=lambda x: x[1],
    reverse=True
)[:30]:
    print(f"{count:5d}  {group}")