from flask import Flask, request, render_template, send_file, flash, redirect, url_for, jsonify, session, make_response
import os
import tempfile
import shutil
try:
    import fitz
    PDF_SUPPORT = True
except ImportError:
    print("WARNING: PyMuPDF not available - PDF processing will not work")
    PDF_SUPPORT = False
    fitz = None
from werkzeug.utils import secure_filename
import zipfile
import time
import threading
from multiprocessing import Pool, cpu_count
import json
from datetime import datetime
from dotenv import load_dotenv
from requests_oauthlib import OAuth2Session
from functools import wraps
import secrets
import base64
import hashlib
from flask_session import Session

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = tempfile.mkdtemp()
app.config['SECRET_KEY'] = os.getenv('SECRET_KEY', 'secret')
app.config['MAX_CONTENT_LENGTH'] = 1000 * 1024 * 1024 * 1024  # 1000TB

# Configure server-side sessions to avoid cookie size limits
app.config['SESSION_TYPE'] = 'filesystem'
app.config['SESSION_FILE_DIR'] = '/tmp/flask_sessions'
app.config['SESSION_PERMANENT'] = False

# Ensure session directory exists and is writable
os.makedirs('/tmp/flask_sessions', exist_ok=True)

try:
    Session(app)
    print("DEBUG: Flask-Session initialized successfully")
except Exception as e:
    print(f"ERROR: Failed to initialize Flask-Session: {e}")
    # Fallback to client-side sessions if server-side fails
    pass

# Load environment variables
load_dotenv()

# Microsoft Entra (Azure AD) Configuration
AZURE_TENANT_ID = os.getenv('AZURE_TENANT_ID')
AZURE_CLIENT_ID = os.getenv('AZURE_CLIENT_ID')
AZURE_CLIENT_SECRET = os.getenv('AZURE_CLIENT_SECRET')
AZURE_AUTHORITY = os.getenv('AZURE_AUTHORITY', f'https://login.microsoftonline.com/{AZURE_TENANT_ID}')
AZURE_REDIRECT_URI = os.getenv('AZURE_REDIRECT_URI', 'http://localhost:5000/auth/callback')
AZURE_SCOPES = os.getenv('AZURE_SCOPES', 'openid profile email').split()

# Validate Azure AD configuration
azure_config_valid = all([
    AZURE_TENANT_ID,
    AZURE_CLIENT_ID, 
    AZURE_CLIENT_SECRET
])

if not azure_config_valid:
    print("Warning: Azure AD configuration is incomplete. Authentication will not work.")
    print("Please check your .env file and ensure all AZURE_* variables are set.")

# OAuth2 configuration
AUTHORIZATION_BASE_URL = f'{AZURE_AUTHORITY}/oauth2/v2.0/authorize'
TOKEN_URL = f'{AZURE_AUTHORITY}/oauth2/v2.0/token'

# Authentication helper functions
def get_oauth_client(state=None):
    """Create OAuth2 client for Azure AD."""
    if not azure_config_valid:
        raise ValueError("Azure AD configuration is incomplete. Please check your .env file.")
    return OAuth2Session(
        client_id=AZURE_CLIENT_ID,
        scope=AZURE_SCOPES,
        state=state,
        redirect_uri=AZURE_REDIRECT_URI
    )

def generate_pkce():
    """Generate PKCE code verifier and challenge."""
    code_verifier = secrets.token_urlsafe(32)
    code_challenge = base64.urlsafe_b64encode(
        hashlib.sha256(code_verifier.encode()).digest()
    ).decode().rstrip('=')
    
    return code_verifier, code_challenge

def login_required(f):
    """Decorator to require authentication for routes."""
    @wraps(f)
    def decorated_function(*args, **kwargs):
        print(f"DEBUG: login_required called for {f.__name__}")
        print(f"DEBUG: azure_config_valid = {azure_config_valid}")
        print(f"DEBUG: session keys = {list(session.keys())}")
        
        if not azure_config_valid:
            # If Azure AD is not configured, show error
            print("DEBUG: Azure AD not configured")
            return "Authentication not configured. Please check your .env file.", 500
        if 'oauth_token' not in session:
            print("DEBUG: No oauth_token in session, redirecting to login")
            return redirect(url_for('login'))
        print("DEBUG: Authentication passed")
        return f(*args, **kwargs)
    return decorated_function

def get_user_info():
    """Get user information from session."""
    if 'user_info' in session:
        return session['user_info']
    return None

# Authentication routes
@app.route('/login')
def login():
    """Initiate Azure AD login flow."""
    if 'oauth_token' in session:
        return redirect(url_for('index'))
    
    # Generate PKCE manually
    code_verifier, code_challenge = generate_pkce()
    session['code_verifier'] = code_verifier
    
    # Build authorization URL manually with PKCE
    from urllib.parse import urlencode
    auth_params = {
        'client_id': AZURE_CLIENT_ID,
        'response_type': 'code',
        'redirect_uri': AZURE_REDIRECT_URI,
        'scope': ' '.join(AZURE_SCOPES),
        'state': secrets.token_urlsafe(16),  # Generate state
        'code_challenge': code_challenge,
        'code_challenge_method': 'S256'
    }
    
    session['oauth_state'] = auth_params['state']
    authorization_url = f"{AUTHORIZATION_BASE_URL}?{urlencode(auth_params)}"
    
    print(f"DEBUG: Authorization URL: {authorization_url}")
    return redirect(authorization_url)

@app.route('/auth/callback')
def auth_callback():
    """Handle Azure AD OAuth2 callback."""
    print(f"DEBUG: Callback URL: {request.url}")
    print(f"DEBUG: Session keys: {list(session.keys())}")
    
    # Check if already authenticated to prevent double callback
    if 'oauth_token' in session:
        print("DEBUG: Already authenticated, redirecting to index")
        response = redirect(url_for('index'))
        response.headers['Cache-Control'] = 'no-cache, no-store, must-revalidate'
        response.headers['Pragma'] = 'no-cache'
        response.headers['Expires'] = '0'
        return response
        response.headers['Pragma'] = 'no-cache'
        response.headers['Expires'] = '0'
        return response
    
    if 'oauth_state' not in session:
        print("DEBUG: No oauth_state in session, redirecting to login")
        response = redirect(url_for('login'))
        response.headers['Cache-Control'] = 'no-cache, no-store, must-revalidate'
        response.headers['Pragma'] = 'no-cache'
        response.headers['Expires'] = '0'
        return response
    
    # Parse the authorization response
    from urllib.parse import parse_qs, urlparse
    parsed_url = urlparse(request.url)
    query_params = parse_qs(parsed_url.query)
    
    print(f"DEBUG: Query params: {query_params}")
    
    if 'error' in query_params:
        print(f"DEBUG: Error from Azure AD: {query_params['error'][0]}")
        return f"Authentication error: {query_params['error'][0]}", 400
    
    if 'code' not in query_params or 'state' not in query_params:
        print("DEBUG: Invalid authorization response - missing code or state")
        return "Invalid authorization response", 400
    
    code = query_params['code'][0]
    state = query_params['state'][0]
    
    if state != session['oauth_state']:
        print(f"DEBUG: State mismatch: {state} != {session['oauth_state']}")
        return "State mismatch", 400
    
    # Exchange code for token using direct HTTP request with PKCE
    import requests
    
    token_data = {
        'client_id': AZURE_CLIENT_ID,
        'client_secret': AZURE_CLIENT_SECRET,
        'code': code,
        'grant_type': 'authorization_code',
        'redirect_uri': AZURE_REDIRECT_URI,
        'scope': ' '.join(AZURE_SCOPES)
    }
    
    # Include code_verifier for PKCE
    if 'code_verifier' in session:
        token_data['code_verifier'] = session['code_verifier']
        print(f"DEBUG: Including code_verifier in token request: {session['code_verifier'][:10]}...")
    
    print(f"DEBUG: Token request data keys: {list(token_data.keys())}")
    
    try:
        token_response = requests.post(TOKEN_URL, data=token_data, timeout=30)
        print(f"DEBUG: Token response status: {token_response.status_code}")
        print(f"DEBUG: Token response: {token_response.text[:200]}...")
        
        if token_response.status_code != 200:
            print(f"DEBUG: Token exchange failed: {token_response.text}")
            return f"Token exchange failed: {token_response.text}", 400
        
        token = token_response.json()
        print(f"DEBUG: Token received successfully")
        session['oauth_token'] = token
        
        # Get user info
        headers = {'Authorization': f"Bearer {token['access_token']}"}
        user_response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers, timeout=30)
        
        if user_response.status_code == 200:
            user_info = user_response.json()
            session['user_info'] = user_info
            print(f"DEBUG: User info retrieved: {user_info.get('displayName', 'Unknown')}")
        else:
            print(f"DEBUG: Failed to get user info: {user_response.text}")
            session['user_info'] = {'displayName': 'Unknown User'}
        
        # Clean up session
        session.pop('oauth_state', None)
        session.pop('code_verifier', None)
        
        print("DEBUG: Authentication successful, redirecting to index")
        response = redirect(url_for('index'))
        response.headers['Cache-Control'] = 'no-cache, no-store, must-revalidate'
        response.headers['Pragma'] = 'no-cache'
        response.headers['Expires'] = '0'
        return response
        
    except requests.exceptions.RequestException as e:
        print(f"DEBUG: Request error during token exchange: {e}")
        return f"Authentication service error: {str(e)}", 500

@app.route('/logout')
def logout():
    """Logout user and clear session."""
    session.clear()
    return redirect(url_for('index'))

# Job storage
JOBS_FILE = 'jobs.json'
RESULTS_DIR = 'results'
os.makedirs(RESULTS_DIR, exist_ok=True)

processing_complete = False
current_output_dir = None

def cleanup_temp_files():
    """Clean up temporary files and directories."""
    import shutil
    import glob
    
    # Clean up any leftover temp directories
    temp_base = tempfile.gettempdir()
    temp_dirs = glob.glob(os.path.join(temp_base, 'tmp*'))
    for temp_dir in temp_dirs:
        try:
            # Only delete if it's a directory and empty or contains our temp files
            if os.path.isdir(temp_dir):
                # Check if it contains our processing files
                contents = os.listdir(temp_dir)
                if any('Addresses' in item or 'Missing_Addresses' in item or 'Originals' in item for item in contents):
                    shutil.rmtree(temp_dir)
                    print(f"Cleaned up temp directory: {temp_dir}")
        except Exception as e:
            print(f"Could not clean up {temp_dir}: {e}")
    
    # Reset global state
    global processing_complete, current_output_dir
    processing_complete = False
    current_output_dir = None
    
    # Clear status file
    if os.path.exists('status.txt'):
        try:
            os.remove('status.txt')
        except:
            pass

def load_jobs():
    """Load jobs from JSON file."""
    if os.path.exists(JOBS_FILE):
        try:
            with open(JOBS_FILE, 'r') as f:
                jobs = json.load(f)
                # Convert timestamp strings back to datetime objects
                for job in jobs.values():
                    if 'timestamp' in job and isinstance(job['timestamp'], str):
                        try:
                            # Handle the format: "2025-11-05 09:10:23.730974"
                            job['timestamp'] = datetime.strptime(job['timestamp'], '%Y-%m-%d %H:%M:%S.%f')
                        except (ValueError, TypeError):
                            # If conversion fails, keep as string
                            pass
                return jobs
        except:
            return {}
    return {}

def save_jobs(jobs):
    """Save jobs to JSON file."""
    with open(JOBS_FILE, 'w') as f:
        json.dump(jobs, f, indent=2, default=str)

@app.route('/', methods=['GET', 'POST'])
@login_required
def index():
    global processing_complete, current_output_dir
    if request.method == 'POST':
        files = request.files.getlist('files')
        if not files or files[0].filename == '':
            flash('No files uploaded')
            return redirect(request.url)
        
        # Clean up any previous processing
        cleanup_temp_files()
        
        # Create temp dirs
        input_dir = tempfile.mkdtemp()
        output_dir = tempfile.mkdtemp()
        current_output_dir = output_dir
        
        for file in files:
            if file and file.filename:
                filename = secure_filename(file.filename)
                if filename.endswith('.zip'):
                    # Extract zip
                    zip_temp = tempfile.mkdtemp()
                    zip_path = os.path.join(zip_temp, filename)
                    file.save(zip_path)
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        zip_ref.extractall(input_dir)
                    shutil.rmtree(zip_temp)
                    print(f"Extracted ZIP: {filename}")
                elif filename.endswith('.pdf'):
                    file.save(os.path.join(input_dir, filename))
                    print(f"Saved PDF: {filename}")
                # For folder upload, files are saved with secure_filename, but since webkitdirectory, filename includes path, but secure_filename flattens it
                else:
                    # Assume folder upload
                    file.save(os.path.join(input_dir, filename))
                    print(f"Saved file from folder: {filename}")
        
        # Create job immediately
        job_start_time = datetime.now()
        job_id = job_start_time.strftime('%Y%m%d_%H%M%S')
        jobs = load_jobs()
        jobs[job_id] = {
            'id': job_id,
            'timestamp': job_start_time,
            'status': 'processing',
            'message': 'Starting processing...',
            'file_count': len([f for f in files if f.filename]),  # Estimate, will be updated
            'progress': 0
        }
        save_jobs(jobs)
        
        # Start processing in background
        processing_complete = False
        thread = threading.Thread(target=process_pdfs_async, args=(input_dir, output_dir, job_id))
        thread.start()
        
        # Redirect to job status page
        return redirect(url_for('job_status', job_id=job_id))
    
    return render_template('index.html', user_info=get_user_info())

@app.route('/status')
def status():
    print("Status requested")
    try:
        with open('status.txt', 'r') as f:
            content = f.read()
        if "Processing complete" in content:
            return jsonify({"status": "complete", "message": content})
        elif "Error" in content:
            return jsonify({"status": "error", "message": content})
        else:
            return jsonify({"status": "processing", "message": content})
    except FileNotFoundError:
        return jsonify({"status": "error", "message": "Processing not complete or file not found"})

@app.route('/results')
@login_required
def results():
    jobs = load_jobs()
    # Sort jobs by timestamp, newest first
    sorted_jobs = sorted(jobs.values(), key=lambda x: x['timestamp'], reverse=True)
    return render_template('results.html', jobs=sorted_jobs, user_info=get_user_info())

@app.route('/jobs')
@login_required
def jobs_page():
    jobs = load_jobs()
    # Sort jobs by timestamp, newest first
    sorted_jobs = sorted(jobs.values(), key=lambda x: x['timestamp'], reverse=True)
    return render_template('jobs.html', jobs=sorted_jobs, user_info=get_user_info())

@app.route('/job/<job_id>')
@login_required
def job_status(job_id):
    jobs = load_jobs()
    if job_id not in jobs:
        flash('Job not found')
        return redirect(url_for('index'))
    
    job = jobs[job_id]
    return render_template('job.html', job=job, user_info=get_user_info())

@app.route('/api/job/<job_id>/status')
def get_job_status(job_id):
    jobs = load_jobs()
    if job_id not in jobs:
        return jsonify({"error": "Job not found"}), 404
    
    job = jobs[job_id]
    return jsonify({
        "id": job["id"],
        "status": job["status"],
        "message": job.get("message", ""),
        "progress": job.get("progress", 0),
        "file_count": job.get("file_count", 0)
    })

def process_pdfs(pdfdir, outputfolder):
    global processing_complete
    status_file = 'status.txt'
    with open(status_file, 'w') as status_f:
        status_f.write("Step 1/4: Finding PDF files...\n")
    print("Step 1/4: Finding PDF files...")
    time.sleep(0.1)
    
    addresses_dir = os.path.join(outputfolder, 'Addresses')
    missing_dir = os.path.join(outputfolder, 'Missing_Addresses')
    os.makedirs(addresses_dir)
    os.makedirs(missing_dir)
    
    with open(os.path.join(addresses_dir, 'Addresses.csv'), 'w') as addr_csv:
        addr_csv.write("File,User,Street,City,State,Zip\n")
    
    with open(os.path.join(missing_dir, 'Missing_Addresses.csv'), 'w') as miss_csv:
        miss_csv.write("File\tUser\n")
    
    pdf_files = []
    for root, dirs, files in os.walk(pdfdir):
        for f in files:
            if f.endswith('.pdf'):
                pdf_files.append(os.path.join(root, f))
    
    msg = f"Step 1/4: Found {len(pdf_files)} PDF files\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print(msg.strip())
    time.sleep(0.1)
    
    # Update job status
    jobs = load_jobs()
    jobs[job_id]['message'] = f"Step 1/4: Found {len(pdf_files)} PDF files"
    jobs[job_id]['file_count'] = len(pdf_files)
    jobs[job_id]['progress'] = 10
    save_jobs(jobs)
    
    msg = "Step 2/4: Processing files...\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print("Step 2/4: Processing files...")
    
    # Update job status
    jobs[job_id]['message'] = "Step 2/4: Processing files..."
    jobs[job_id]['progress'] = 20
    save_jobs(jobs)
    
    time.sleep(0.1)
    
    # Process files in parallel
    with Pool(processes=cpu_count()) as pool:
        results = pool.map(process_single_pdf, pdf_files)
    
    for result in results:
        f = result['filename']
        if result['type'] == 'missing':
            # Missing address
            shutil.copy(result['filepath'], os.path.join(missing_dir, f))
            with open(os.path.join(missing_dir, 'Missing_Addresses.csv'), 'a') as miss_csv:
                miss_csv.write(f"{f}\t{result['user']}\n")
            msg = f"  -> Moved to Missing_Addresses\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
        elif result['type'] == 'address':
            # Valid address
            shutil.copy(result['filepath'], os.path.join(addresses_dir, f))
            with open(os.path.join(addresses_dir, 'Addresses.csv'), 'a') as addr_csv:
                addr_csv.write(f"{f},{result['user']},{result['address']},{result['city']},{result['state']},{result['zip']}\n")
            msg = f"  -> Moved to Addresses (User: {result['user']}, Address: {result['address']}, {result['city']})\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
        elif result['type'] == 'error':
            msg = f"  -> Error processing {f}: {result['error']}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
        time.sleep(0.05)
    
    msg = "Step 2/4: Sorting complete. Starting merging...\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print(msg.strip())
    
    # Update job status
    jobs[job_id]['message'] = "Step 2/4: Sorting complete. Starting merging..."
    jobs[job_id]['progress'] = 80
    save_jobs(jobs)
    
    time.sleep(0.1)
    
    # Merge by page count
    page_groups = {}
    for filepath in pdf_files:
        f = os.path.basename(filepath)
        doc = fitz.open(filepath)
        pages = len(doc)
        doc.close()
        if pages not in page_groups:
            page_groups[pages] = []
        page_groups[pages].append(filepath)
    
    msg = f"Step 3/4: Found {len(page_groups)} different page counts to merge\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print(msg.strip())
    time.sleep(0.1)
    
    for pages, files in page_groups.items():
        msg = f"Merging {len(files)} files with {pages} pages each\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        time.sleep(0.05)
        try:
            merger = fitz.open()
            for file in files:
                doc = fitz.open(file)
                merger.insert_pdf(doc)
                doc.close()
            output_path = os.path.join(outputfolder, f'Statement-Pages-{pages}.pdf')
            merger.save(output_path)
            merger.close()
            msg = f"  -> Created {output_path}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
        except Exception as e:
            msg = f"  -> Error merging files with {pages} pages: {str(e)}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
        time.sleep(0.05)
    
    msg = f"Processing complete! Processed {len(pdf_files)} files total\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print(msg.strip())
    
    # Update job status
    jobs[job_id]['message'] = f"Processing complete! Processed {len(pdf_files)} files total"
    jobs[job_id]['progress'] = 90
    save_jobs(jobs)
    
    time.sleep(0.1)
    
    # Step 4/4: Creating ZIP file
    msg = "Step 4/4: Creating ZIP file...\n"
    with open(status_file, 'a') as status_f:
        status_f.write(msg)
    print(msg.strip())
    
    # Update job status
    jobs[job_id]['message'] = "Step 4/4: Creating ZIP file..."
    jobs[job_id]['progress'] = 95
    save_jobs(jobs)
    
    time.sleep(0.1)
    
    # Create zip with unique name
    zip_filename = f'{job_id}.zip'
    zip_path = os.path.join(RESULTS_DIR, zip_filename)
    
    try:
        # Create zip with only merged PDF documents (no CSV files)
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            # Add merged statement PDFs from root output folder
            for file in os.listdir(outputfolder):
                if file.endswith('.pdf') and file.startswith('Statement-Pages-'):
                    file_path = os.path.join(outputfolder, file)
                    zipf.write(file_path, file)
        
        # Save job metadata
        jobs = load_jobs()
        jobs[job_id]['status'] = 'completed'
        jobs[job_id]['zip_path'] = zip_path
        jobs[job_id]['zip_filename'] = zip_filename
        jobs[job_id]['message'] = f'Processing complete! ZIP created: {zip_filename}'
        jobs[job_id]['progress'] = 100
        save_jobs(jobs)
        
        msg = f"ZIP created: {zip_filename}\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        
    except Exception as e:
        msg = f"Error creating ZIP: {str(e)}\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        # Still save job as failed
        jobs = load_jobs()
        jobs[job_id]['status'] = 'failed'
        jobs[job_id]['error'] = str(e)
        jobs[job_id]['message'] = f'Error creating ZIP: {str(e)}'
        save_jobs(jobs)
    
    processing_complete = True
    global current_output_dir
    current_output_dir = outputfolder

def process_pdfs_async(pdfdir, outputfolder, job_id):
    global processing_complete
    try:
        jobs = load_jobs()
        if job_id not in jobs:
            return
        
        status_file = 'status.txt'
        with open(status_file, 'w') as status_f:
            status_f.write("Step 1/4: Finding PDF files...\n")
        print("Step 1/4: Finding PDF files...")
        
        # Update job status
        jobs[job_id]['message'] = "Step 1/4: Finding PDF files..."
        jobs[job_id]['progress'] = 5
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        addresses_dir = os.path.join(outputfolder, 'Addresses')
        missing_dir = os.path.join(outputfolder, 'Missing_Addresses')
        os.makedirs(addresses_dir)
        os.makedirs(missing_dir)
        
        pdf_files = []
        for root, dirs, files in os.walk(pdfdir):
            for f in files:
                if f.endswith('.pdf'):
                    pdf_files.append(os.path.join(root, f))
        
        msg = f"Step 1/4: Found {len(pdf_files)} PDF files\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        
        # Update job status
        jobs = load_jobs()
        jobs[job_id]['message'] = f"Step 1/4: Found {len(pdf_files)} PDF files"
        jobs[job_id]['file_count'] = len(pdf_files)
        jobs[job_id]['progress'] = 10
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        msg = "Step 2/4: Processing files...\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print("Step 2/4: Processing files...")
        
        # Update job status
        jobs[job_id]['message'] = "Step 2/4: Processing files..."
        jobs[job_id]['progress'] = 20
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        # Parallel processing
        num_processes = min(cpu_count(), len(pdf_files))
        with Pool(processes=num_processes) as pool:
            results = pool.map(process_single_pdf, pdf_files)
        
        # Process results sequentially
        with open(os.path.join(addresses_dir, 'Addresses.csv'), 'w') as addr_csv:
            addr_csv.write("File,User,Street,City,State,Zip\n")
        
        with open(os.path.join(missing_dir, 'Missing_Addresses.csv'), 'w') as miss_csv:
            miss_csv.write("File\tUser\n")
        
        page_groups = {}
        for i, result in enumerate(results, 1):
            progress_percent = 20 + int((i / len(results)) * 60)  # 20-80% range for processing
            
            msg = f"Processing result {i}/{len(results)}: {result.get('filename', 'unknown')}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
            
            # Update job status periodically
            if i % 100 == 0 or i == len(results):  # Update every 100 files or at the end
                jobs[job_id]['message'] = f"Processing result {i}/{len(results)}: {result.get('filename', 'unknown')}"
                jobs[job_id]['progress'] = progress_percent
                save_jobs(jobs)
            
            if result['type'] == 'address':
                with open(os.path.join(addresses_dir, 'Addresses.csv'), 'a') as addr_csv:
                    addr_csv.write(f"{result['filename']},{result['user']},{result['address']},{result['city']},{result['state']},{result['zip']}\n")
                shutil.copy(result['filepath'], os.path.join(addresses_dir, result['filename']))
                msg = f"  -> Moved to Addresses\n"
                with open(status_file, 'a') as status_f:
                    status_f.write(msg)
                print(msg.strip())
            elif result['type'] == 'missing':
                with open(os.path.join(missing_dir, 'Missing_Addresses.csv'), 'a') as miss_csv:
                    miss_csv.write(f"{result['filename']}\t{result['user']}\n")
                shutil.copy(result['filepath'], os.path.join(missing_dir, result['filename']))
                msg = f"  -> Moved to Missing_Addresses\n"
                with open(status_file, 'a') as status_f:
                    status_f.write(msg)
                print(msg.strip())
            elif result['type'] == 'error':
                msg = f"  -> Error: {result['error']}\n"
                with open(status_file, 'a') as status_f:
                    status_f.write(msg)
                print(msg.strip())
            
            # Group for merging
            if 'pages' in result:
                pages = result['pages']
                if pages not in page_groups:
                    page_groups[pages] = []
                page_groups[pages].append(result['filepath'])
            
            time.sleep(0.01)  # Small delay for status updates
        
        msg = "Step 2/4: Sorting complete. Starting merging...\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        
        # Update job status
        jobs[job_id]['message'] = "Step 2/4: Sorting complete. Starting merging..."
        jobs[job_id]['progress'] = 80
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        # Merge by page count
        msg = f"Step 3/4: Found {len(page_groups)} different page counts to merge\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        time.sleep(0.1)
        
        for pages, files in page_groups.items():
            msg = f"Merging {len(files)} files with {pages} pages each\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
            time.sleep(0.05)
            try:
                merger = fitz.open()
                for file in files:
                    doc = fitz.open(file)
                    merger.insert_pdf(doc)
                    doc.close()
                output_path = os.path.join(outputfolder, f'Statement-Pages-{pages}.pdf')
                merger.save(output_path)
                merger.close()
                msg = f"  -> Created {output_path}\n"
                with open(status_file, 'a') as status_f:
                    status_f.write(msg)
                print(msg.strip())
            except Exception as e:
                msg = f"  -> Error merging files with {pages} pages: {str(e)}\n"
                with open(status_file, 'a') as status_f:
                    status_f.write(msg)
                print(msg.strip())
            time.sleep(0.05)
        
        msg = f"Processing complete! Processed {len(pdf_files)} files total\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        
        # Update job status
        jobs[job_id]['message'] = f"Processing complete! Processed {len(pdf_files)} files total"
        jobs[job_id]['progress'] = 90
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        # Step 4/4: Creating ZIP file
        msg = "Step 4/4: Creating ZIP file...\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        
        # Update job status
        jobs[job_id]['message'] = "Step 4/4: Creating ZIP file..."
        jobs[job_id]['progress'] = 95
        save_jobs(jobs)
        
        time.sleep(0.1)
        
        # Create zip with unique name
        zip_filename = f'{job_id}.zip'
        zip_path = os.path.join(RESULTS_DIR, zip_filename)
        
        try:
            # Create zip with only merged PDF documents (no CSV files)
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                # Add merged statement PDFs from root output folder
                for file in os.listdir(outputfolder):
                    if file.endswith('.pdf') and file.startswith('Statement-Pages-'):
                        file_path = os.path.join(outputfolder, file)
                        zipf.write(file_path, file)
        
            # Save job metadata
            jobs = load_jobs()
            jobs[job_id]['status'] = 'completed'
            jobs[job_id]['zip_path'] = zip_path
            jobs[job_id]['zip_filename'] = zip_filename
            jobs[job_id]['message'] = f'Processing complete! ZIP created: {zip_filename}'
            jobs[job_id]['progress'] = 100
            save_jobs(jobs)
            
            msg = f"ZIP created: {zip_filename}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
            
        except Exception as e:
            msg = f"Error creating ZIP: {str(e)}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
            # Still save job as failed
            jobs = load_jobs()
            jobs[job_id]['status'] = 'failed'
            jobs[job_id]['error'] = str(e)
            jobs[job_id]['message'] = f'Error creating ZIP: {str(e)}'
            save_jobs(jobs)
    
    except Exception as e:
        msg = f"Processing error: {e}"
        with open('status.txt', 'w') as f:
            f.write(msg)
        # Mark job as failed
        jobs = load_jobs()
        if job_id in jobs:
            jobs[job_id]['status'] = 'failed'
            jobs[job_id]['error'] = str(e)
            jobs[job_id]['message'] = msg
            save_jobs(jobs)
        processing_complete = True  # Still mark as complete to allow retry

@app.route('/download/<job_id>')
@login_required
def download_job(job_id):
    jobs = load_jobs()
    if job_id in jobs and jobs[job_id]['status'] == 'completed':
        zip_path = jobs[job_id]['zip_path']
        if os.path.exists(zip_path):
            return send_file(zip_path, as_attachment=True, download_name=jobs[job_id]['zip_filename'])
    return jsonify({"error": "Job not found or not completed"}), 404

@app.route('/download')
def download():
    # Legacy route - redirect to results page
    return redirect(url_for('results'))

@app.route('/delete/<job_id>', methods=['POST'])
@login_required
def delete_job(job_id):
    jobs = load_jobs()
    if job_id in jobs:
        job = jobs[job_id]
        # Delete ZIP file if it exists
        if 'zip_path' in job and os.path.exists(job['zip_path']):
            try:
                os.remove(job['zip_path'])
            except:
                pass  # Ignore errors
        # Remove from jobs
        del jobs[job_id]
        save_jobs(jobs)
        flash(f'Job {job_id} deleted successfully')
    else:
        flash('Job not found')
    return redirect(url_for('results'))

@app.route('/clear', methods=['POST'])
def clear_processing():
    try:
        cleanup_temp_files()
        return jsonify({"message": "Temporary files cleared and processing state reset."})
    except Exception as e:
        return jsonify({"message": f"Error clearing files: {str(e)}"}), 500

@app.route('/clear_and_upload')
def clear_and_upload():
    cleanup_temp_files()
    flash('Processing state cleared. Ready to upload new files.')
    return redirect(url_for('index'))

def process_single_pdf(filepath):
    """Process a single PDF file and return the result."""
    try:
        f = os.path.basename(filepath)
        doc = fitz.open(filepath)
        text = ""
        for page in doc:
            text += page.get_text()
        pages = len(doc)
        doc.close()
        
        lines = text.split('\n')
        if len(lines) < 15:
            user = ""
            addr1 = ""
            addr2 = ""
        else:
            user = lines[11].strip()
            addr1 = lines[12].strip()
            addr2 = lines[13].strip()
        
        if not addr1:
            return {'type': 'missing', 'filename': f, 'user': user, 'filepath': filepath, 'pages': pages}
        else:
            if ',' in addr2:
                city, rest = addr2.split(',', 1)
                city = city.strip()
                rest_parts = rest.strip().split()
                state = rest_parts[0] if rest_parts else ''
                zip_code = ' '.join(rest_parts[1:]) if len(rest_parts) > 1 else ''
            else:
                city = addr2
                state = ''
                zip_code = ''
            return {'type': 'address', 'filename': f, 'user': user, 'address': addr1, 'city': city, 'state': state, 'zip': zip_code, 'filepath': filepath, 'pages': pages}
    except Exception as e:
        return {'type': 'error', 'filename': os.path.basename(filepath), 'error': str(e), 'filepath': filepath}

if __name__ == '__main__':
    try:
        from waitress import serve
        print("Starting with Waitress server (supports large uploads)...")
        serve(app, host='0.0.0.0', port=5000, max_request_body_size=1000*1024*1024*1024)  # 1TB
    except ImportError:
        print("Waitress not available, using Flask dev server...")
        app.run(host='0.0.0.0', port=5000, debug=True)