#!/usr/bin/env python3
"""Handle partial or corrupted RDF downloads by extracting usable data."""

import gzip
import sys
from pathlib import Path

from rich.console import Console

console = Console()


def extract_partial_gz(gz_file: Path, output_file: Path, max_errors: int = 10):
    """Extract as much as possible from a potentially corrupted gzip file.

    Args:
        gz_file: Path to the gzip file
        output_file: Path to output the extracted data
        max_errors: Maximum number of consecutive errors before stopping

    Returns:
        Number of lines extracted
    """
    lines_extracted = 0
    consecutive_errors = 0

    try:
        with (
            gzip.open(gz_file, "rt", encoding="utf-8", errors="ignore") as f_in,
            open(output_file, "w", encoding="utf-8") as f_out,
        ):
            while consecutive_errors < max_errors:
                    try:
                        line = f_in.readline()
                        if not line:
                            break
                        f_out.write(line)
                        lines_extracted += 1
                        consecutive_errors = 0  # Reset error counter on success

                        # Show progress every 100k lines
                        if lines_extracted % 100000 == 0:
                            console.print(
                                f"[dim]Extracted {lines_extracted:,} lines...[/dim]"
                            )
                    except Exception:
                        consecutive_errors += 1
                        if consecutive_errors >= max_errors:
                            console.print(
                                f"[yellow]Stopping after {consecutive_errors} "
                                f"consecutive errors[/yellow]"
                            )
                            break
    except EOFError:
        console.print("[yellow]Reached end of compressed data[/yellow]")
    except Exception as e:
        console.print(f"[red]Error opening file: {e}[/red]")

    return lines_extracted


def main():
    if len(sys.argv) != 3:
        console.print("Usage: python handle_partial_downloads.py <input.gz> <output>")
        sys.exit(1)

    input_file = Path(sys.argv[1])
    output_file = Path(sys.argv[2])

    if not input_file.exists():
        console.print(f"[red]Input file not found: {input_file}[/red]")
        sys.exit(1)

    console.print(f"[cyan]Extracting partial data from: {input_file}[/cyan]")
    console.print(f"[cyan]Output: {output_file}[/cyan]")

    lines = extract_partial_gz(input_file, output_file)

    if lines > 0:
        output_size_mb = output_file.stat().st_size / (1024 * 1024)
        console.print(
            f"[green]✓ Successfully extracted {lines:,} lines "
            f"({output_size_mb:.2f} MB)[/green]"
        )
    else:
        console.print("[red]No data could be extracted[/red]")
        sys.exit(1)


if __name__ == "__main__":
    main()