import sys
import camelot
import polars as pl
import signal
import argparse
from rich.console import Console
from rich.progress import track

console = Console()

class PDFTableParser:
    def __init__(self, input_files, output_files, delimiter, edge_tol, row_tol, pages):
        self.input_files = input_files
        self.output_files = output_files
        self.delimiter = delimiter
        self.edge_tol = edge_tol
        self.row_tol = row_tol
        self.pages = pages

    def read_tables(self, file_name):
        try:
            console.print(f"Reading tables from {file_name}...")
            tables = camelot.read_pdf(file_name, flavor='stream', edge_tol=self.edge_tol, row_tol=self.row_tol, pages=self.pages)
            console.print(f"Found {len(tables)} tables in {file_name}.")
            return tables
        except Exception as e:
            console.print(f"[red]Error reading {file_name}: {e}[/red]")
            return None

    def save_tables_as_csv(self, tables, output_file):
        try:
            console.print(f"Saving tables to {output_file}...")
            df = pl.concat([pl.DataFrame(table.df) for table in tables])
            df.write_csv(output_file, separator=self.delimiter)
            console.print(f"Saved tables to {output_file}.")
        except Exception as e:
            console.print(f"[red]Error saving to {output_file}: {e}[/red]")

    def estimate_processing_time(self, file_name):
        try:
            with open(file_name, 'rb') as f:
                content = f.read().decode('utf-8', errors='ignore')
            pages = content.count('\n')
            words = len(content.split())
            chars = len(content)
            estimated_time = (lines / 1000) + (words / 1000) + (chars / 1000)
            console.print(f"Estimated processing time for {file_name}: {estimated_time:.2f} seconds.")
            return estimated_time
        except Exception as e:
            console.print(f"[red]Error estimating processing time for {file_name}: {e}[/red]")
            return 0

    def process_files(self):
        for input_file, output_file in track(zip(self.input_files, self.output_files), description="Processing files"):
            self.estimate_processing_time(input_file)
            tables = self.read_tables(input_file)
            if tables:
                self.save_tables_as_csv(tables, output_file)

def handle_signal(signum, frame):
    console.print("\n[red]Process interrupted.[/red]")
    sys.exit(1)

if __name__ == "__main__":
    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    parser = argparse.ArgumentParser(description="PDF Table Parser")
    parser.add_argument("input_files", nargs='+', help="List of input PDF files")
    parser.add_argument("output_files", nargs='+', help="List of output CSV files")
    parser.add_argument("--delimiter", default=',', help="Output file delimiter (default: ,)")
    parser.add_argument("--edge_tol", type=int, default=50, help="Tolerance parameter used to specify the distance between text and table edges (default: 50)")
    parser.add_argument("--row_tol", type=int, default=50, help="Tolerance parameter used to specify the distance between table rows (default: 50)")
    parser.add_argument("--pages", type=str, default='all', help="Pages you can pass the number of pages to process. (default: all)")
    parser.add_argument("--webui", action='store_true', help="Launch the web UI")

    args = parser.parse_args()

    if len(args.input_files) != len(args.output_files):
        console.print("[red]The number of input files and output files must match.[/red]")
        sys.exit(1)

    if args.webui:
        webui = WebUI()
        webui.run()
    else:
        main(args)