import argparse
import enum
import logging
import math
import pathlib
import re
import subprocess
import sys
import urllib.request

from typing import (
	ClassVar,
	Optional,
)

logger = logging.getLogger(__name__)


class Command(enum.Enum):
	list_installed = 'list-installed'
	compile = 'compile'
	download = 'download'
	archive = 'archive'


class parse_rate_t:
	class constants_t:
		rate_re: ClassVar[re.Pattern[str]] = re.compile(r'^(\d+(?:\.\d+)?)\s*([bBkKmMgGpPtT]?)(?:[iI]?[bB])?(?:/s)?$')

		units: ClassVar[dict[str, int]] = {
			'': 0,
			'b': 0,
			'B': 0,
			'k': 1,
			'K': 1,
			'm': 2,
			'M': 2,
			'g': 3,
			'G': 3,
			't': 4,
			'T': 4,
			'p': 5,
			'P': 5,
		}

	@staticmethod
	def parse(s: str) -> int:
		m = parse_rate_t.constants_t.rate_re.match(s.strip())
		if not m:
			raise ValueError('invalid rate: %s' % s)

		value = float(m.group(1))
		unit = m.group(2)

		power = parse_rate_t.constants_t.units.get(unit, 0)

		return int(value * (1024**power))


class downloader_t:
	class constants_t:
		class backend_t(enum.Enum):
			urllib = 'urllib'
			curl = 'curl'
			aria2c = 'aria2c'

	@staticmethod
	def download(
		url: str,
		dest: pathlib.Path,
		backend: 'downloader_t.constants_t.backend_t',
		limit_rate: int,
	) -> None:
		dest.parent.mkdir(parents=True, exist_ok=True)

		if backend is downloader_t.constants_t.backend_t.urllib:
			urllib.request.urlretrieve(url, str(dest))
		elif backend is downloader_t.constants_t.backend_t.curl:
			cmd = [
				'curl',
				'-fSL',
				'--limit-rate',
				'%d' % limit_rate,
				'-o',
				str(dest),
				url,
			]
			subprocess.check_call(cmd)
		elif backend is downloader_t.constants_t.backend_t.aria2c:
			cmd = [
				'aria2c',
				'--max-download-limit=%d' % limit_rate,
				'-d',
				str(dest.parent),
				'-o',
				dest.name,
				url,
			]
			subprocess.check_call(cmd)
		else:
			raise NotImplementedError


class download_requirements_t:
	@staticmethod
	def parse_requirements(txt: str) -> list[tuple[str, str]]:
		entries: list[tuple[str, str]] = []
		url: Optional[str] = None

		for line in txt.splitlines():
			line = line.strip()
			if line == '':
				continue
			if line.startswith('#'):
				candidate = line[1:].strip()
				if '/' in candidate and '://' in candidate:
					url = candidate
				continue

			parts = line.split()
			if len(parts) == 0:
				continue

			pkg_spec = parts[0]

			if url is not None:
				filename = url.rsplit('/', 1)[-1] if '/' in url else pkg_spec
				entries.append((url, filename))
				url = None

		return entries


def _find_cached_pkg(
	cache_dir: pathlib.Path,
	name: str,
	version: str,
) -> Optional[pathlib.Path]:
	"""Find a cached .pkg.tar.* file for a given package name and version."""
	for suffix in ['.pkg.tar.zst', '.pkg.tar.xz', '.pkg.tar.gz', '.pkg.tar.bz2', '.pkg.tar']:
		for arch in ['x86_64', 'any']:
			candidate = cache_dir / ('%s-%s-%s%s' % (name, version, arch, suffix))
			if candidate.exists():
				return candidate
	return None


def main(argv: Optional[list[str]] = None) -> int:
	if argv is None:
		argv = sys.argv[1:]

	logging.basicConfig(level=logging.INFO)

	parser = argparse.ArgumentParser(
		prog='online-fxreader-pr34-archlinux',
		description='Arch Linux package management tools',
	)
	parser.add_argument(
		'command',
		choices=[o.value for o in Command],
	)

	options, args = parser.parse_known_args(argv)
	options.command = Command(options.command)

	if options.command is Command.list_installed:
		import hashlib

		from .pacman import pacman_t

		list_parser = argparse.ArgumentParser()
		list_parser.add_argument(
			'--format',
			choices=['plain', 'constraints', 'compiled'],
			default='plain',
			help='plain: name version; constraints: name>=version; compiled: name==version with optional hashes',
		)
		list_parser.add_argument(
			'--generate-hashes',
			action='store_true',
			default=False,
			help='include sha256 from local /var/cache/pacman/pkg/ files; fails if file not found for any package',
		)
		list_parser.add_argument(
			'--db-path',
			dest='db_path',
			default='/var/lib/pacman',
			help='pacman db path, default /var/lib/pacman',
		)
		list_parser.add_argument(
			'--pkg-cache-dir',
			dest='pkg_cache_dir',
			default='/var/cache/pacman/pkg',
			help='local pacman package cache directory, default /var/cache/pacman/pkg',
		)

		list_options = list_parser.parse_args(args)

		installed = pacman_t.list_installed_simple(
			db_path=pathlib.Path(list_options.db_path),
		)

		pkg_cache_dir = pathlib.Path(list_options.pkg_cache_dir)

		if list_options.format == 'plain':
			for name, version in installed:
				print('%s %s' % (name, version))
		elif list_options.format == 'constraints':
			for name, version in installed:
				print('%s>=%s' % (name, version))
		elif list_options.format == 'compiled':
			missing_hashes: list[str] = []

			for name, version in installed:
				line = '%s==%s' % (name, version)

				if list_options.generate_hashes:
					pkg_file = _find_cached_pkg(
						pkg_cache_dir,
						name,
						version,
					)

					if pkg_file is not None:
						h = hashlib.sha256()
						with open(pkg_file, 'rb') as fh:
							while True:
								chunk = fh.read(65536)
								if not chunk:
									break
								h.update(chunk)
						line += ' --hash=sha256:%s' % h.hexdigest()
					else:
						missing_hashes.append(name)

				print(line)

			if len(missing_hashes) > 0:
				logger.error(
					"can't determine checksum of installed package(s) - no cached file found for %d package(s): %s" % (len(missing_hashes), missing_hashes)
				)
				return 1

		return 0
	elif options.command is Command.compile:
		compile_parser = argparse.ArgumentParser()
		compile_parser.add_argument(
			'packages',
			nargs='*',
		)
		compile_parser.add_argument(
			'-r',
			dest='requirements_file',
			default=None,
			help='path to file with package constraints (one per line)',
		)
		compile_parser.add_argument(
			'--index',
			dest='index_url',
			default=None,
			help='mirror URL',
		)
		compile_parser.add_argument(
			'--archive-date',
			dest='archive_date',
			default=None,
			help='Arch Linux Archive date (e.g. 2024/01/15)',
		)
		compile_parser.add_argument(
			'--offline',
			action='store_true',
			default=False,
		)
		compile_parser.add_argument(
			'--no-cache',
			action='store_true',
			default=False,
		)
		compile_parser.add_argument(
			'--generate-hashes',
			action='store_true',
			default=False,
		)
		compile_parser.add_argument(
			'--cache-dir',
			dest='cache_dir',
			default=None,
		)
		compile_parser.add_argument(
			'--repos',
			nargs='*',
			default=['core', 'extra', 'multilib'],
		)
		compile_parser.add_argument(
			'--arch',
			default='x86_64',
		)
		compile_parser.add_argument(
			'--backend',
			choices=['python', 'solv'],
			default='solv',
		)
		compile_parser.add_argument(
			'--archive-cache',
			dest='archive_cache',
			default=None,
			help='path to archive cache dir (with archlinux_cache.db from archive sync); loads all synced dates into the solver pool',
		)
		compile_parser.add_argument(
			'--reference',
			default=None,
			help='path to previously compiled requirements file to use as version pins',
		)
		compile_parser.add_argument(
			'--resolution-strategy',
			dest='resolution_strategy',
			choices=['upgrade-all', 'pin-referenced'],
			default='upgrade-all',
			help='upgrade-all: resolve fresh; pin-referenced: keep referenced versions, only upgrade explicitly requested packages',
		)

		compile_options = compile_parser.parse_args(args)

		from .models import compile_options_t, resolution_strategy_t

		packages: list[str] = list(compile_options.packages)

		if compile_options.requirements_file is not None:
			for line in pathlib.Path(compile_options.requirements_file).read_text().splitlines():
				line = line.strip()
				if line != '' and not line.startswith('#'):
					packages.append(line)

		opts = compile_options_t(
			packages=packages,
			index_url=compile_options.index_url,
			archive_date=compile_options.archive_date,
			offline=compile_options.offline,
			no_cache=compile_options.no_cache,
			generate_hashes=compile_options.generate_hashes,
			repos=compile_options.repos,
			arch=compile_options.arch,
			cache_dir=compile_options.cache_dir,
			reference=compile_options.reference,
			resolution_strategy=resolution_strategy_t(compile_options.resolution_strategy),
		)

		try:
			if compile_options.backend == 'solv':
				from .solv_backend import compile_solv_t, repo_store_t

				stores = None
				if compile_options.archive_cache is not None:
					from .cache_db import cache_db_t

					archive_cache_dir = pathlib.Path(compile_options.archive_cache)
					db_path = archive_cache_dir / 'archlinux_cache.db'
					if db_path.exists():
						cache_db = cache_db_t(db_path)
						indices = cache_db.load_all_indices()
						cache_db.close()
						stores = [repo_store_t(index=idx) for idx in indices]

				result = compile_solv_t.compile(opts, stores=stores)
			else:
				from .compile import compile_t

				result = compile_t.compile(opts)
		except RuntimeError as e:
			logger.error(str(e))
			return 1

		print(result.txt)

		return 0
	elif options.command is Command.download:
		download_parser = argparse.ArgumentParser()
		download_parser.add_argument(
			'-r',
			dest='requirements',
			required=True,
			help='path to compiled requirements file',
		)
		download_parser.add_argument(
			'-d',
			dest='dest_dir',
			required=True,
			help='destination directory for downloaded packages',
		)
		download_parser.add_argument(
			'--downloader',
			choices=[o.value for o in downloader_t.constants_t.backend_t],
			default='urllib',
		)
		download_parser.add_argument(
			'--limit-rate',
			dest='limit_rate',
			default='128KiB/s',
			help='download speed limit (e.g. 128KiB/s, 1MiB/s, 512K), default 128KiB/s',
		)

		download_options = download_parser.parse_args(args)

		dest_dir = pathlib.Path(download_options.dest_dir)
		dest_dir.mkdir(parents=True, exist_ok=True)

		backend = downloader_t.constants_t.backend_t(download_options.downloader)
		limit_rate = parse_rate_t.parse(download_options.limit_rate)

		requirements_txt = pathlib.Path(download_options.requirements).read_text()
		entries = download_requirements_t.parse_requirements(requirements_txt)

		count = 0
		for url, filename in entries:
			dest_path = dest_dir / filename

			if dest_path.exists():
				logger.info(dict(msg='already downloaded', path=str(dest_path)))
			else:
				logger.info(dict(msg='downloading', url=url, dest=str(dest_path), backend=backend.value, limit_rate=limit_rate))
				downloader_t.download(
					url=url,
					dest=dest_path,
					backend=backend,
					limit_rate=limit_rate,
				)

			count += 1

		logger.info(dict(msg='download complete', count=count))

		return 0
	elif options.command is Command.archive:
		from . import archive as _archive

		return _archive.main(args)
	else:
		raise NotImplementedError


if __name__ == '__main__':
	sys.exit(main())
