"""
This module provides functions for creating and managing archives.

- repair_zip_file: Try to repair a corrupted zip file by ignoring some of the errors
- make_archive: Create a zip archive from a source directory with consistent file timestamps.
- archive_cli: Main entry point for command line usage

.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/archive_module.gif
  :alt: stouputils archive examples
"""

# pyright: reportUnusedVariable=false
# Imports
import fnmatch
import os
from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo

from .decorators import LogLevels, handle_error
from .io import clean_path, super_copy
from .print import debug, error, info


# Function that repair a corrupted zip file (ignoring some of the errors)
@handle_error()
def repair_zip_file(file_path: str, destination: str) -> bool:
	""" Try to repair a corrupted zip file by ignoring some of the errors

	This function manually parses the ZIP file structure to extract files
	even when the ZIP file is corrupted. It reads the central directory
	entries and attempts to decompress each file individually.

	Args:
		file_path		(str):	Path of the zip file to repair
		destination		(str):	Destination of the new file
	Returns:
		bool: Always returns True unless any strong error

	Examples:

	.. code-block:: python

		> repair_zip_file("/path/to/source.zip", "/path/to/destination.zip")
	"""
	# Check
	if not os.path.exists(file_path):
		raise FileNotFoundError(f"File '{file_path}' not found")
	dirname: str = os.path.dirname(destination)
	if dirname and not os.path.exists(dirname):
		raise FileNotFoundError(f"Directory '{dirname}' not found")

	import struct
	import zlib

	# Read the entire ZIP file into memory
	with open(file_path, 'rb') as f:
		data = f.read()

	# Find central directory entries
	CENTRAL_SIG = b'PK\x01\x02'
	entries: list[dict[str, int | str]] = []
	idx = 0

	while True:
		idx = data.find(CENTRAL_SIG, idx)
		if idx == -1:
			break
		# Ensure enough length for central directory header
		if idx + 46 > len(data):
			break

		header = data[idx:idx+46]
		try:
			(
				sig,
				ver_made, ver_needed, flags, comp_method, mtime, mdate,
				crc, csize, usize,
				name_len, extra_len, comm_len,
				disk_start, int_attr,
				ext_attr, local_off
			) = struct.unpack('<4s6H3L3H2H2L', header)

			name_start = idx + 46
			if name_start + name_len > len(data):
				idx += 4
				continue

			name = data[name_start:name_start+name_len].decode('utf-8', errors='replace')
			entries.append({
				'name': name,
				'comp_method': comp_method,
				'csize': csize,
				'usize': usize,
				'local_offset': local_off,
				'crc': crc
			})
		except (struct.error, UnicodeDecodeError):
			# Skip corrupted entries
			pass

		idx += 4

	# Create a new ZIP file with recovered entries
	with ZipFile(destination, "w", compression=ZIP_DEFLATED) as new_zip_file:
		for entry in entries:
			try:
				# Get the local header to find data start
				lo: int = int(entry['local_offset'])
				if lo + 30 > len(data):
					continue

				lh = data[lo:lo+30]
				try:
					_, _, _, _, _, _, _, _, _, name_len, extra_len = struct.unpack('<4sHHHHHLLLHH', lh)
				except struct.error:
					continue

				data_start: int = lo + 30 + name_len + extra_len
				if data_start + int(entry['csize']) > len(data):
					continue

				comp_data = data[data_start:data_start+int(entry['csize'])]

				# Decompress the data
				try:
					if int(entry['comp_method']) == 0:  # No compression
						content = comp_data[:int(entry['usize'])]
					elif int(entry['comp_method']) == 8:  # Deflate compression
						content = zlib.decompress(comp_data, -zlib.MAX_WBITS)
					else:
						# Unsupported compression method, skip
						continue

					# Write to new ZIP file
					new_zip_file.writestr(str(entry['name']), content)

				except (zlib.error, Exception):
					# If decompression fails, try to write raw data as a fallback
					try:
						new_zip_file.writestr(f"{entry['name']!s}.corrupted", comp_data)
					except Exception:
						# Skip completely corrupted entries
						continue

			except Exception:
				# Skip any entries that cause errors
				continue

	return True

# Function that makes an archive with consistency (same zip file each time)
@handle_error()
def make_archive(
	source: str,
	destinations: list[str] | str | None = None,
	override_time: None | tuple[int, int, int, int, int, int] = None,
	create_dir: bool = False,
	ignore_patterns: str | None = None,
) -> bool:
	""" Create a zip archive from a source directory with consistent file timestamps.
	(Meaning deterministic zip file each time)

	Creates a zip archive from the source directory and copies it to one or more destinations.
	The archive will have consistent file timestamps across runs if override_time is specified.
	Uses maximum compression level (9) with ZIP_DEFLATED algorithm.

	Args:
		source				(str):						The source folder to archive
		destinations		(list[str]|str):			The destination folder(s) or file(s) to copy the archive to
		override_time		(None | tuple[int, ...]):	The constant time to use for the archive
			(e.g. (2024, 1, 1, 0, 0, 0) for 2024-01-01 00:00:00)
		create_dir			(bool):						Whether to create the destination directory if it doesn't exist
		ignore_patterns		(str | None):				Glob pattern(s) to ignore files. Can be a single pattern or comma-separated patterns (e.g. "*.pyc" or "*.pyc,__pycache__,*.log")
	Returns:
		bool: Always returns True unless any strong error
	Examples:

	.. code-block:: python

		> make_archive("/path/to/source", "/path/to/destination.zip")
		> make_archive("/path/to/source", ["/path/to/destination.zip", "/path/to/destination2.zip"])
		> make_archive("src", "hello_from_year_2085.zip", override_time=(2085,1,1,0,0,0))
		> make_archive("src", "output.zip", ignore_patterns="*.pyc")
		> make_archive("src", "output.zip", ignore_patterns="__pycache__")
		> make_archive("src", "output.zip", ignore_patterns="*.pyc,__pycache__,*.log")
	"""
	# Fix copy_destinations type if needed
	if destinations is None:
		destinations = []
	if destinations and isinstance(destinations, str):
		destinations = [destinations]
	if not destinations:
		raise ValueError("destinations must be a list of at least one destination")

	# Create directories if needed
	if create_dir:
		for dest_file in destinations:
			dest_file = clean_path(dest_file)
			parent_dir = os.path.dirname(dest_file)
			if parent_dir and not os.path.exists(parent_dir):
				os.makedirs(parent_dir, exist_ok=True)

	# Create the archive
	destination: str = clean_path(destinations[0])
	destination = destination if ".zip" in destination else destination + ".zip"

	# Parse ignore patterns (can be a single pattern or comma-separated patterns)
	ignore_pattern_list: list[str] = []
	if ignore_patterns:
		ignore_pattern_list = [pattern.strip() for pattern in ignore_patterns.split(',')]

	def should_ignore(path: str) -> bool:
		"""Check if a file or directory path should be ignored based on patterns."""
		if not ignore_pattern_list:
			return False
		for pattern in ignore_pattern_list:
			if fnmatch.fnmatch(os.path.basename(path), pattern) or fnmatch.fnmatch(path, pattern):
				return True
		return False

	with ZipFile(destination, "w", compression=ZIP_DEFLATED, compresslevel=9) as zip:
		for root, dirs, files in os.walk(source):
			# Filter out ignored directories in-place to prevent walking into them
			dirs[:] = [d for d in dirs if not should_ignore(d)]

			for file in files:
				file_path: str = clean_path(os.path.join(root, file))
				rel_path = os.path.relpath(file_path, source)

				# Skip files that match any ignore pattern
				if should_ignore(file) or should_ignore(rel_path):
					continue

				info: ZipInfo = ZipInfo(rel_path)
				info.compress_type = ZIP_DEFLATED
				if override_time:
					info.date_time = override_time
				with open(file_path, "rb") as f:
					zip.writestr(info, f.read())

	# Copy the archive to the destination(s)
	for dest_file in destinations[1:]:
		@handle_error(Exception, message=f"Unable to copy '{destination}' to '{dest_file}'", error_log=LogLevels.WARNING)
		def internal(src: str, dest: str) -> None:
			super_copy(src, dest, create_dir=create_dir)
		internal(destination, clean_path(dest_file))

	return True


# Main entry point for command line usage
def archive_cli() -> None:
	""" Main entry point for command line usage.

	Examples:

	.. code-block:: bash

		# Repair a corrupted zip file
		python -m stouputils.archive repair /path/to/corrupted.zip /path/to/repaired.zip

		# Create a zip archive
		python -m stouputils.archive make /path/to/source /path/to/destination.zip

		# Create a zip archive with ignore patterns
		python -m stouputils.archive make /path/to/source /path/to/destination.zip --ignore "*.pyc,__pycache__"
	"""
	import argparse
	import sys

	parser = argparse.ArgumentParser(description="Archive utilities")
	subparsers = parser.add_subparsers(dest="command", help="Available commands")

	# Repair command
	repair_parser = subparsers.add_parser("repair", help="Repair a corrupted zip file")
	repair_parser.add_argument("input_file", help="Path to the corrupted zip file")
	repair_parser.add_argument("output_file", nargs="?", help="Path to the repaired zip file (optional, defaults to input_file with '_repaired' suffix)")

	# Make archive command
	archive_parser = subparsers.add_parser("make", help="Create a zip archive")
	archive_parser.add_argument("source", help="Source directory to archive")
	archive_parser.add_argument("destination", help="Destination zip file")
	archive_parser.add_argument("--ignore", help="Glob patterns to ignore (comma-separated)")
	archive_parser.add_argument("--create-dir", action="store_true", help="Create destination directory if it doesn't exist")

	args = parser.parse_args()

	if args.command == "repair":
		input_file = args.input_file
		if args.output_file:
			output_file = args.output_file
		else:
			# Generate default output filename
			base, ext = os.path.splitext(input_file)
			output_file = f"{base}_repaired{ext}"

		debug(f"Repairing '{input_file}' to '{output_file}'...")
		try:
			repair_zip_file(input_file, output_file)
			info(f"Successfully repaired zip file: {output_file}")
		except Exception as e:
			error(f"Error repairing zip file: {e}", exit=False)
			sys.exit(1)

	elif args.command == "make":
		debug(f"Creating archive from '{args.source}' to '{args.destination}'...")
		try:
			make_archive(
				source=args.source,
				destinations=args.destination,
				create_dir=args.create_dir,
				ignore_patterns=args.ignore
			)
			info(f"Successfully created archive: {args.destination}")
		except Exception as e:
			error(f"Error creating archive: {e}", exit=False)
			sys.exit(1)

	else:
		parser.print_help()
		sys.exit(1)


if __name__ == "__main__":
	archive_cli()


