#!/usr/bin/env python3

# Copyright (C) 2018-2023 Simon Quigley <tsimonq2@ubuntu.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import json
import libcalamares
from time import strftime
import urllib.request
from urllib.error import HTTPError, URLError
import socket
import logging
import distro
import xml.etree.ElementTree as ET
from os import remove


SUBDOMAINS_BY_COUNTRY_CODE = {
    "US": "us.", "AU": "au.", "SE": "no.", "NO": "no.",
    "NZ": "nz.", "NL": "nl.", "KR": "kr.", "DE": "de.",
    "GE": "ge.", "PF": "pf.", "CZ": "cz.", "HR": "hr."
}


def get_sources_template():
    return """# Automatically generated by Calamares on {date}.
# See http://help.ubuntu.com/community/UpgradeNotes for how to upgrade to
# newer versions of {distro}.

Types: deb
URIs: {url}
Suites: {codename} {codename}-updates
Components: main universe
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
"""


def get_sec_sources_template():
    return """Types: deb
URIs: http://security.ubuntu.com/ubuntu
Suites: {codename}-security
Components: main universe
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
"""


def get_restricted_sources_template():
    return """Types: deb
URIs: {url}
Suites: {codename} {codename}-updates
Components: restricted multiverse
Signed-By: /usr/share/keyrings/ubuntu-archive-keyring.gpg
Enabled: no
"""


def get_country_code():
    if not libcalamares.globalstorage.value("hasInternet"):
        return ""
    
    geoip_config = libcalamares.job.configuration["geoip"]
    
    try:
        with urllib.request.urlopen(geoip_config["url"], timeout=75) as resp:
            if geoip_config["style"] == "json":
                return json.loads(resp.read().decode())["country_code"]
            elif geoip_config["style"] == "xml":
                return ET.parse(resp).getroot().find("CountryCode").text
    except (HTTPError, URLError, socket.timeout):
        logging.error("Failed to get country code.")
    return ""


def get_subdomain_by_country(country_code):
    return SUBDOMAINS_BY_COUNTRY_CODE.get(country_code, "")


def write_file(path, content):
    with open(path, "w") as f:
        f.write(content)


def run():
    country_code = get_country_code()
    subdomain = get_subdomain_by_country(country_code)
    base_url = "http://{}{}/ubuntu".format(subdomain, libcalamares.job.configuration["baseUrl"])
    codename = distro.codename()

    root_mount_point = libcalamares.globalstorage.value("rootMountPoint")

    sources = get_sources_template().format(date=strftime("%Y-%m-%d"), distro=libcalamares.job.configuration["distribution"], url=base_url, codename=codename)
    write_file(f"{root_mount_point}/etc/apt/sources.list.d/ubuntu.sources", sources)
    
    sec_sources = get_sec_sources_template().format(codename=codename)
    write_file(f"{root_mount_point}/etc/apt/sources.list.d/ubuntu_security.sources", sec_sources)

    restricted_sources = get_restricted_sources_template().format(url=base_url, codename=codename)
    write_file(f"{root_mount_point}/etc/apt/sources.list.d/ubuntu_nonfree.sources", restricted_sources)
    
    remove(f"{root_mount_point}/etc/apt/sources.list")

    libcalamares.globalstorage.insert("mirrorURL", base_url)
    libcalamares.globalstorage.insert("ubuntuCodename", codename)