#!/usr/bin/env python3
import argparse
import urllib.request
import json
import signal
import re
import shutil
import textwrap
from concurrent.futures import ThreadPoolExecutor, as_completed
from urllib.error import HTTPError, URLError

# Color codes
COLORS = {
    'reset': '\033[0m',
    'bold': '\033[1m',
    'red': '\033[91m',
    'green': '\033[92m',
    'yellow': '\033[93m',
    'header': '\033[94m'
}

def colorize(text, color):
    """Add ANSI color codes if output is a terminal"""
    if not hasattr(colorize, 'is_tty'):
        colorize.is_tty = __import__('sys').stdout.isatty()
    return f"{COLORS[color]}{text}{COLORS['reset']}" if colorize.is_tty else text

# def get_terminal_width(default=80):
#     """Get terminal width with fallback"""
#     try:
#         return shutil.get_terminal_size().columns
#     except:
#         return default

signal.signal(signal.SIGINT, lambda s, f: exit(1))

REQUEST_HEADERS = {
    'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) checkpkg/1.0',
    'Accept': 'application/json'
}

PM_MAPPINGS = {
    'apt': ['debian_', 'ubuntu_'],
    'dnf': ['fedora_'],
    'pacman': ['arch'],
    'aur': ['aur'],
    'apk': ['alpine_'],
    'zypper': ['opensuse_']
}

REPO_FORMATS = {
    'debian_': "Debian {}",
    'ubuntu_': "Ubuntu {}",
    'fedora_': "Fedora {}",
    'arch': "Arch",
    'aur': "AUR",
    'alpine_': "Alpine {}",
    'opensuse_': "openSUSE {}"
}

def version_key(version):
    """Create a sorting key for version comparison"""
    return [
        (0, int(part)) if part.isdigit() else (1, part.lower())
        for part in re.findall(r'(\d+|\D+)', version)
    ]

def get_package_manager(repo):
    for pm, patterns in PM_MAPPINGS.items():
        if any(repo.startswith(p) for p in patterns):
            return pm
    return None

def format_repository(repo):
    for pattern, fmt in REPO_FORMATS.items():
        if repo.startswith(pattern):
            parts = repo.split('_')
            return fmt.format(parts[1] if len(parts) > 1 else '')
    return repo

def fetch_package_data(package):
    try:
        req = urllib.request.Request(
            f'https://repology.org/api/v1/project/{package}',
            headers=REQUEST_HEADERS
        )
        with urllib.request.urlopen(req, timeout=10) as response:
            return json.load(response)
    except HTTPError as e:
        if e.code == 403:
            print(colorize(f"Error: Repology blocked the request for {package} (try again later)", 'red'))
        return None
    except Exception:
        return None

def print_table(output, headers, selected_pms, args):
    """Print formatted table with consistent alignment or list for small terminals"""
    # Terminal dimensions
    terminal_width = shutil.get_terminal_size().columns
    
    # Calculate package availability for coloring
    pkg_availability = {}
    for pkg in args.packages:
        versions = output.get(pkg, {})
        available_count = sum(1 for pm in selected_pms if versions.get(pm, 'Not found') != 'Not found')
        
        if available_count == 0:
            pkg_availability[pkg] = 'red'  # Not available anywhere
        elif available_count == len(selected_pms):
            pkg_availability[pkg] = 'green'  # Available everywhere
        else:
            pkg_availability[pkg] = 'yellow'  # Available in some places
    
    # Determine minimum required widths
    min_pkg_width = max(len(pkg) for pkg in args.packages) + 2
    min_pm_width = 10
    min_required_width = min_pkg_width + (min_pm_width * len(selected_pms)) + (3 * len(selected_pms))
    
    # If terminal is too narrow for the table, use list format instead
    if terminal_width < min_required_width and len(selected_pms) > 1:
        print_list_format(output, headers, selected_pms, args, pkg_availability)
        return
    
    # Calculate column widths
    padding = 1  # Space on each side of content
    
    # Package column width
    pkg_col_width = min(min_pkg_width + 4, max(min_pkg_width, terminal_width // (len(selected_pms) + 3)))
    
    # PM column widths (divide remaining space equally)
    remaining_width = terminal_width - pkg_col_width - (3 * len(selected_pms))
    pm_col_width = max(min_pm_width, remaining_width // len(selected_pms))
    
    col_widths = [pkg_col_width] + [pm_col_width] * len(selected_pms)
    
    # Print header row
    header_cells = []
    for i, header in enumerate(headers):
        text = header.center(col_widths[i] - (2 * padding))
        cell = " " * padding + colorize(text, 'header') + " " * padding
        header_cells.append(cell)
    
    print(" | ".join(header_cells))
    
    # Print separator line
    total_width = sum(col_widths) + (3 * (len(col_widths) - 1))
    print("-" * total_width)
    
    # Print each package row
    for pkg_idx, pkg in enumerate(args.packages):
        versions = output.get(pkg, {})
        
        # First collect all data for this package
        package_data = []
        
        # Package name (first column)
        package_data.append([colorize(pkg, pkg_availability[pkg])])
        
        # Version data for each package manager
        for pm in selected_pms:
            version = versions.get(pm, 'Not found')
            
            if version == 'Not found':
                package_data.append([colorize('-', 'red')])
                continue
            
            # Extract version number and repositories
            version_parts = []
            match = re.match(r'(.*?)\s+\((.*)\)$', version)
            
            if match:
                ver_num, repos = match.groups()
                version_parts.append(colorize(ver_num, 'green'))
                
                # Format repositories
                repo_lines = []
                repo_text = "(" + repos + ")"
                
                # Wrap repository text if needed
                avail_width = col_widths[len(package_data)] - (2 * padding)
                if len(repo_text) <= avail_width:
                    repo_lines.append(colorize(repo_text, 'green'))
                else:
                    # Handle wrapping for repositories
                    repo_parts = repos.split(', ')
                    current_line = "("
                    
                    for repo in repo_parts:
                        if len(current_line) + len(repo) + 2 <= avail_width:
                            if current_line != "(":
                                current_line += ", "
                            current_line += repo
                        else:
                            if current_line != "(":
                                current_line += ")"
                                repo_lines.append(colorize(current_line, 'green'))
                            current_line = "  " + repo
                    
                    if current_line != "(":
                        current_line += ")" if not current_line.startswith("  ") else ""
                        repo_lines.append(colorize(current_line, 'green'))
                
                # Combined version and repo lines
                package_data.append([version_parts[0]] + repo_lines)
            else:
                # Simple version string
                package_data.append([colorize(version, 'green')])
        
        # Determine max number of lines needed
        max_lines = max(len(column) for column in package_data)
        
        # Print all lines for this package
        for line_idx in range(max_lines):
            row_cells = []
            
            for col_idx, col_data in enumerate(package_data):
                if line_idx < len(col_data):
                    # Actual content
                    content = col_data[line_idx]
                    content_plain = re.sub(r'\033\[\d+m', '', content)
                    
                    # Calculate padding
                    left_pad = padding
                    right_pad = max(0, col_widths[col_idx] - len(content_plain) - left_pad)
                    
                    cell = " " * left_pad + content + " " * right_pad
                else:
                    # Empty cell
                    cell = " " * col_widths[col_idx]
                
                row_cells.append(cell)
            
            print(" | ".join(row_cells))
        
        # Add separator between packages
        if pkg_idx < len(args.packages) - 1:
            print("·" * total_width)

def print_list_format(output, headers, selected_pms, args, pkg_availability):
    """Print packages in a vertical list format for narrow terminals"""
    terminal_width = shutil.get_terminal_size().columns
    
    for pkg_idx, pkg in enumerate(args.packages):
        pkg_color = pkg_availability[pkg]
        versions = output.get(pkg, {})
        
        # Print package header with color based on availability
        print(f"\n{colorize('Package:', 'bold')} {colorize(pkg, pkg_color)}")
        print("-" * min(40, terminal_width - 2))
        
        # Print versions for each package manager
        for pm in selected_pms:
            pm_name = headers[selected_pms.index(pm) + 1]  # Get friendly display name
            version = versions.get(pm, 'Not found')
            
            if version == 'Not found':
                print(f"{colorize(pm_name, 'header')}: {colorize('-', 'red')}")
            else:
                # Extract version and repo information
                match = re.match(r'(.*?)\s+\((.*)\)$', version)
                if match:
                    ver_num, repos = match.groups()
                    
                    # Handle long repository lists with wrapping
                    if len(pm_name) + len(ver_num) + len(repos) + 5 > terminal_width:
                        print(f"{colorize(pm_name, 'header')}: {colorize(ver_num, 'green')}")
                        # Wrap repositories with proper indentation
                        wrapper = textwrap.TextWrapper(
                            width=terminal_width - 4,
                            initial_indent="    ",
                            subsequent_indent="    "
                        )
                        wrapped = wrapper.fill(f"({repos})")
                        print(colorize(wrapped, 'green'))
                    else:
                        print(f"{colorize(pm_name, 'header')}: {colorize(ver_num, 'green')} ({repos})")
                else:
                    print(f"{colorize(pm_name, 'header')}: {colorize(version, 'green')}")
        
        # Add separator between packages
        if pkg_idx < len(args.packages) - 1:
            print("\n" + "-" * min(40, terminal_width - 2))

def main():
    parser = argparse.ArgumentParser(description='Package search tool')
    parser.add_argument('--all', action='store_true')
    parser.add_argument('--apt', action='store_true')
    parser.add_argument('--dnf', action='store_true')
    parser.add_argument('--pacman', action='store_true')
    parser.add_argument('--apk', action='store_true')
    parser.add_argument('--zypper', action='store_true')
    parser.add_argument('--aur', action='store_true')
    parser.add_argument('packages', nargs='+')
    args = parser.parse_args()

    selected_pms = [pm for pm, flag in [
        ('apt', args.apt or args.all),
        ('dnf', args.dnf or args.all),
        ('pacman', args.pacman or args.all),
        ('apk', args.apk or args.all),
        ('zypper', args.zypper or args.all),
        ('aur', args.aur or args.all)
    ] if flag]

    if not selected_pms:
        print(colorize("Error: No package managers selected", 'red'))
        return

    results = {}
    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = {executor.submit(fetch_package_data, pkg): pkg for pkg in args.packages}
        for future in as_completed(futures):
            pkg = futures[future]
            try:
                data = future.result()
                results[pkg] = data or []
            except Exception as e:
                print(colorize(f"Error processing {pkg}: {str(e)}", 'red'))
                results[pkg] = []

    output = {}
    for pkg, entries in results.items():
        pm_versions = {pm: {'version': '', 'repos': set(), 'key': []} for pm in selected_pms}
        
        for entry in entries:
            repo = entry.get('repo', '')
            version = entry.get('version', 'N/A')
            pm = get_package_manager(repo)
            
            if pm in selected_pms and version != 'N/A':
                repo_fmt = format_repository(repo)
                current_key = version_key(version)
                stored = pm_versions[pm]
                
                if not stored['key'] or current_key > stored['key']:
                    stored['version'] = version
                    stored['repos'] = {repo_fmt}
                    stored['key'] = current_key
                elif current_key == stored['key']:
                    stored['repos'].add(repo_fmt)
        
        output[pkg] = {}
        for pm in selected_pms:
            data = pm_versions[pm]
            if data['version']:
                repos = ', '.join(sorted(data['repos']))
                output[pkg][pm] = f"{data['version']} ({repos})"
            else:
                output[pkg][pm] = 'Not found'

    print_table(output, ['Package'] + selected_pms, selected_pms, args)

if __name__ == '__main__':
    main()