#!/usr/bin/env python3
"""
Classes to handle apt:
    FilteredSourcesList: Used by Apt class - creates a list with debian and solydxk source entries.
    SolydAptSources: main apt sources class

Author: Arjen Balfoort
Date: 2026-01-16
Dependencies: aptitude, dpkg, apt-mark
"""

import sys
import re
import subprocess
import os
import gettext
import json
import tempfile

from typing import Optional
from enum import StrEnum
from utils import (
    get_domain_from_uri,
    get_debian_version,
    show_dialog,
    DialogIcon,
    SudoCommandThread,
    DebSoftwareApp
)

# https://apt-team.pages.debian.net/python-apt/library/aptsources.sourceslist.html
# Only support the deb822 .sources format
try:
    from aptsources.sourceslist import SourcesList, Deb822SourceEntry
except ImportError:
    sys.exit()

import gi
gi.require_version('Gtk', '4.0')
from gi.repository import Gtk

_ = gettext.translation('solydxk-system', fallback=True).gettext

class Repositories(StrEnum):
    """ Possible repostitories """
    SOLYDXK = "solydxk.com"
    DEBIAN = "debian.org"
    DEBIAN_OTHER = "other.debian.org"
    OTHER = "other"


class SolydAptSources():
    """Apt class"""
    def __init__(self, app: DebSoftwareApp):
        # Called property
        self.distro_comps: dict = {'debian':
                                ['main',
                                'contrib',
                                'non-free',
                                'non-free-firmware'],
                            'solydxk':
                                ['main',
                                'upstream',
                                'import']
                            }

        self.keyrings: dict = {'debian': 'debian-archive-keyring.gpg',
                         'solydxk': 'solydxk-archive-keyring.gpg'}

        self.app: DebSoftwareApp = app
        self.debian_mirrors: dict = self.load_mirrors_from_json('debian')
        self.solydxk_mirrors: dict = self.load_mirrors_from_json('solydxk')
        self.debian_releases: dict = self.load_debian_releases_from_json()
        self.sources: SourcesList = SourcesList(deb822=True)
        self.solydxk_sources: list[Deb822SourceEntry] = []
        self.debian_sources: list[Deb822SourceEntry] = []
        self.other_sources: list[Deb822SourceEntry] = []
        self.architectures: list[str] = self._get_architectures()

        self._fill_repo_sources()
        self._check_sources()

    def _check_sources(self):
        """
        Check if the sources files are missing and create them if necessary 
        """
        # You can only get here if the user removed the debian or solydxk source entries
        entries_data = []

        if not self.solydxk_sources:
            deb_version = get_debian_version()
            if deb_version > 0:
                first_key = next(iter(self.solydxk_mirrors))
                entry = self.make_new_source_entry(types=['deb'],
                                                   uris=[self.solydxk_mirrors[first_key][0]["uri"]],
                                                   suites=[f'solydxk-{deb_version}'],
                                                   file='/etc/apt/sources.list.d/solydxk.sources')
                # Save the section and file path for the JSON file
                source_section = self.get_source_section(entry)
                entries_data.append({
                    "section": source_section,
                    "file": entry.file
                })

        if not self.debian_sources:
            # Get the correct release name
            deb_version = get_debian_version()
            try:
                debian_release = self.debian_releases[deb_version].lower()
            except KeyError:
                debian_release = 'stable'

            first_key = next(iter(self.debian_mirrors))
            deb_entry = self.make_new_source_entry(types=['deb'],
                                        uris=[self.debian_mirrors[first_key][0]["uri"]],
                                        suites=[debian_release, f'{debian_release}-updates'],
                                        file='/etc/apt/sources.list.d/debian.sources')

            sec_entry = self.make_new_source_entry(types=['deb'],
                                        uris=["https://security.debian.org/debian-security/"],
                                        suites=[f'{debian_release}-security'])
            # Save the section and file path for the JSON file
            source_section = self.get_source_section(deb_entry) + '\n' + \
                             self.get_source_section(sec_entry)
            entries_data.append({
                "section": source_section,
                "file": deb_entry.file
            })

        if entries_data:
            res = show_dialog(title=_("Repositories"),
                              text=_("Some source files are missing.\n"
                                     "Do you want to create them now?\n\n"
                                     "This will take a few seconds."),
                              main_window=self.app.window,
                              icon=DialogIcon.QUESTION,
                              buttons=["Yes", "No"])
            if res == Gtk.ResponseType.YES:
                # Write to a temporary JSON file
                with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                    json.dump(entries_data, f)
                    json_file = f.name

                # pkexec will prompt for root password via polkit
                command = ["pkexec",
                           "--disable-internal-agent",
                           "python3",
                           os.path.join(self.app.script_dir, "repositories_sudo.py"),
                           json_file]

                thread = SudoCommandThread(command=command,
                                           subject=_("Repositories"),
                                           logging_handler=self.app.logging)
                thread.start()
                thread.join()

                # Update the sources
                self.sources = SourcesList(deb822=True)
                self._fill_repo_sources()


    def _fill_repo_sources(self) -> None:
        """ Fill the sources repository lists """
        for source_entry in self.sources.list:
            if self.is_solydxk(source_entry):
                self.solydxk_sources.append(source_entry)
            elif self.is_debian(source_entry):
                # Check for other debian repositories
                domain = get_domain_from_uri(source_entry.uri)
                if domain.startswith("deb.") or domain.startswith("ftp."):
                    self.debian_sources.append(source_entry)
                else:
                    self.other_sources.append(source_entry)
            else:
                self.other_sources.append(source_entry)

    def load_mirrors_from_json(self, distribution: str) -> dict:
        """
        Load mirrors from json file
        """
        file_path = os.path.join(self.app.data_dir, 'debian_mirrors.json')
        if distribution == 'solydxk':
            file_path = os.path.join(self.app.data_dir, 'solydxk_mirrors.json')

        if not os.path.exists(file_path):
            return {}

        with open(file=file_path, mode='r', encoding='utf-8') as f:
            return json.load(f)

    def load_debian_releases_from_json(self) -> dict:
        """
        Load Debian releases from json file
        """
        file_path = os.path.join(self.app.data_dir, 'debian_releases.json')

        if not os.path.exists(file_path):
            return {}

        with open(file=file_path, mode='r', encoding='utf-8') as f:
            loaded = json.load(f)
            # Convert keys to integers
            return {int(k): v for k, v in loaded.items()}

    def _get_architectures(self)-> list:
        """
        Get all supported architectures
        Used as default if no architectures are provided
        """
        architectures = []
        try:
            result = subprocess.run(
                ["dpkg", "--print-architecture"],
                capture_output=True,
                text=True,
                check=True
            )
            architectures = result.stdout.strip().split('\n')
        except subprocess.CalledProcessError as e:
            self.app.logging.error("Error: %s.", e)

        try:
            result = subprocess.run(
                ["dpkg", "--print-foreign-architectures"],
                capture_output=True,
                text=True,
                check=True
            )
            architectures += result.stdout.strip().split('\n')
        except subprocess.CalledProcessError as e:
            self.app.logging.error("Error: %s.", e)

        self.app.logging.debug("Architectures: %s", ', '.join(architectures))
        return architectures

    def get_signed_by(self, source_entry: Deb822SourceEntry) -> str:
        """
        Get the path to the keyring file

        Args:
            source_entry (Deb822SourceEntry): Source entry

        Returns:
            str: Path to keyring
        """
        keyring = ''
        match = re.search(pattern=r'signed-by[=:\s]+([^\]\n]+)',
                          string=source_entry.line,
                          flags=re.IGNORECASE)
        if match:
            # Keyring can be moved: check if it still exists
            if os.path.exists(match.group(1)):
                keyring = match.group(1)
                self.app.logging.info("Keyring still exists: %s", keyring)

        # Keyring not found: check if installed on system
        if not keyring:
            for repo in Repositories:
                if repo in source_entry.line:
                    keyring = self._search_keyring_path(repo)
                    self.app.logging.info("Found keyring: %s", keyring)
        return keyring

    def _search_keyring_path(self, search_string: str) -> Optional[str]:
        """
        Get the path to the keyring file
        Called by get_signed_by()

        Args:
            repo (str): Repository name

        Returns:
            str: Path to keyring
        """
        if not search_string:
            return None

        keyring_name = search_string
        if 'solydxk' in search_string:
            keyring_name = self.keyrings['solydxk']
        else:
            keyring_name = self.keyrings['debian']
        self.app.logging.debug("Keyring name: %s", keyring_name)

        try:
            result = subprocess.run(
                ["ls", "/usr/share/keyrings"],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            ).stdout.splitlines()
        except subprocess.CalledProcessError:
            return None

        for keyring in result:
            if keyring_name in keyring:
                return f'/usr/share/keyrings/{keyring}'

        return None

    # Called function (was new_apt_source)
    def make_new_source_entry(self,
                              types: list[str],
                              uris: list[str],
                              suites: list[str],
                              comps: list[str]=None,
                              architectures: list[str]=None,
                              signed_by: str=None,
                              disabled: bool=False,
                              comment: str=None,
                              file: str=None) -> Deb822SourceEntry:
        """Create apt source dictionary

        Args:
            types (list[str]): apt source types
            uri (str): apt source uri
            suites (list[str]): apt source suites
            comps (list[str]): apt source components
            architectures (list[str], optional): apt source arhitectures. Defaults to None.
            signed_by (str, optional): path to keyring file. Defaults to None.
                If not provided it is chosen by repo (_get_signed_by function).
            disabled (bool, optional): sets whether or not the apt source is disabled.
                Defaults to False.
            comment (str, optional): apt source comment. Defaults to ''.
            file (str, optional): apt sources save file. Defaults to None.
                If not provided a source file is chosen by repo.

        Returns:
            Deb822SourceEntry
        """
        # Check current used system
        if file is None:
            for source_entry in self.sources.list:
                for repository in Repositories:
                    if repository in source_entry.uri and \
                       repository in uris:
                        file=source_entry.file
                        break
                if file is not None:
                    break

        # Get the deb822 section string
        archs = architectures if architectures is not None else self.architectures
        signed_by_path = signed_by if signed_by is not None else self._search_keyring_path(uris[0])

        if comps is None:
            if Repositories.SOLYDXK.value in uris[0]:
                comps = self.distro_comps['solydxk']
            else:
                comps = self.distro_comps['debian']

        # Set default comment for this repo
        source_comment = f"# {suites[0]} repository" if not comment else comment

        # Enabled string
        enabled = 'no' if disabled else 'yes'

        source_section = f"{source_comment}\n" \
        f"Enabled: {enabled}\n" \
        f"Types: {' '.join(types)}\n" \
        f"URIs: {' '.join(uris)}\n" \
        f"Suites: {' '.join(suites)}\n" \
        f"Components: {' '.join(comps)}\n" \
        f"Architectures: {' '.join(archs)}\n" \
        f"Signed-By: {signed_by_path}\n\n"

        return Deb822SourceEntry(section=source_section, file=file)

    def get_source_section(self, source_entry: Deb822SourceEntry) -> str:
        source_comment = source_entry.comment.strip()
        source_comment = f"{get_domain_from_uri(source_entry.uri)} repository" \
                         if not source_comment else source_comment
        archs = source_entry.architectures if source_entry.architectures else self.architectures
        enabled = 'no' if source_entry.disabled else 'yes'
        signed_by_path = self.get_signed_by(source_entry)

        return f"# {source_comment}\n" \
               f"Enabled: {enabled}\n" \
               f"Types: {' '.join(source_entry.types)}\n" \
               f"URIs: {' '.join(source_entry.uris)}\n" \
               f"Suites: {' '.join(source_entry.suites)}\n" \
               f"Components: {' '.join(source_entry.comps)}\n" \
               f"Architectures: {' '.join(archs)}\n" \
               f"Signed-By: {signed_by_path}\n"

    def is_debian(self, source_entry: Deb822SourceEntry) -> bool:
        """
        Check if a source entry is a Debian repository.

        Args:
            source_entry (Deb822SourceEntry): The source entry to check.

        Returns:    
            bool: True if the source entry is a Debian repository, False otherwise.
        """
        if Repositories.DEBIAN.value in source_entry.uri:
            return True

        for _country, mirror_list in self.debian_mirrors.items():
            for mirror in mirror_list:
                if mirror.get('domain') in source_entry.uri:
                    return True
        return False

    def is_solydxk(self, source_entry: Deb822SourceEntry) -> bool:
        """
        Check if a source entry is a Solydxk repository.

        Args:
            source_entry (Deb822SourceEntry): The source entry to check.

        Returns:    
            bool: True if the source entry is a Solydxk repository, False otherwise.
        """
        if Repositories.SOLYDXK.value in source_entry.uri:
            return True

        for _country, mirror_list in self.solydxk_mirrors.items():
            for mirror in mirror_list:
                if mirror.get('domain') in source_entry.uri:
                    return True
        return False
