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
import logging

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

LOGO_B64 = None  # Use static file instead of base64 to avoid URL length issues

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['SECRET_KEY'] = os.getenv('SECRET_KEY', 'secret')
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024 * 1024  # 100GB (reduced from 1000TB for practicality)

# Configure server-side sessions to avoid cookie size limits
app.config['SESSION_TYPE'] = 'filesystem'
app.config['SESSION_FILE_DIR'] = os.path.abspath(os.path.join(os.path.dirname(__file__), 'sessions'))
app.config['SESSION_PERMANENT'] = False

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

# Ensure uploads directory exists
os.makedirs(os.path.abspath(os.path.join(os.path.dirname(__file__), 'uploads')), 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()

# Startup diagnostics
logger.info("=" * 60)
logger.info("APPLICATION STARTUP DIAGNOSTICS")
logger.info("=" * 60)
logger.info(f"Working directory: {os.getcwd()}")
logger.info(f"App directory: {os.path.dirname(os.path.abspath(__file__))}")
logger.info(f"PDF Support (PyMuPDF): {PDF_SUPPORT}")
logger.info("=" * 60)

@app.context_processor
def inject_logo():
    # Return static file path instead of base64
    return {'logo_url': '/static/trbc-logo-full.png'}

# 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 = os.path.abspath(os.path.join(os.path.dirname(__file__), 'jobs.json'))
RESULTS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), 'results'))
os.makedirs(RESULTS_DIR, exist_ok=True)

# Check directory permissions at startup
app_dir = os.path.dirname(os.path.abspath(__file__))
uploads_dir = os.path.join(app_dir, 'uploads')
sessions_dir = os.path.join(app_dir, 'sessions')

print(f"\nDirectory Checks:")
print(f"  Uploads: {uploads_dir}")
print(f"    Exists: {os.path.exists(uploads_dir)}")
print(f"    Writable: {os.access(uploads_dir, os.W_OK) if os.path.exists(uploads_dir) else 'N/A'}")
print(f"  Sessions: {sessions_dir}")
print(f"    Exists: {os.path.exists(sessions_dir)}")
print(f"    Writable: {os.access(sessions_dir, os.W_OK) if os.path.exists(sessions_dir) else 'N/A'}")
print(f"  Results: {RESULTS_DIR}")
print(f"    Exists: {os.path.exists(RESULTS_DIR)}")
print(f"    Writable: {os.access(RESULTS_DIR, os.W_OK) if os.path.exists(RESULTS_DIR) else 'N/A'}")
print(f"  Jobs file: {JOBS_FILE}")
print(f"    Exists: {os.path.exists(JOBS_FILE)}")
if os.path.exists(JOBS_FILE):
    print(f"    Writable: {os.access(JOBS_FILE, os.W_OK)}")
print("=" * 60)

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."""
    print(f"Loading jobs from {JOBS_FILE}")
    print(f"Jobs file exists: {os.path.exists(JOBS_FILE)}")
    if os.path.exists(JOBS_FILE):
        try:
            with open(JOBS_FILE, 'r') as f:
                jobs = json.load(f)
                print(f"Loaded raw jobs: {list(jobs.keys())}")
                # Convert timestamp strings back to datetime objects
                for job_id, job in jobs.items():
                    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')
                            print(f"Converted timestamp for job {job_id}: {job['timestamp']}")
                        except (ValueError, TypeError) as e:
                            print(f"Failed to convert timestamp for job {job_id}: {job['timestamp']} - {e}")
                            # If conversion fails, keep as string
                            pass
                    else:
                        print(f"Job {job_id} has timestamp type: {type(job.get('timestamp'))} value: {job.get('timestamp')}")
                print(f"Processed {len(jobs)} jobs")
                return jobs
        except Exception as e:
            print(f"Error loading jobs: {e}")
            return {}
    else:
        print("Jobs file does not exist")
    return {}

def save_jobs(jobs):
    """Save jobs to JSON file."""
    print(f"Saving {len(jobs)} jobs to {JOBS_FILE}")
    try:
        with open(JOBS_FILE, 'w') as f:
            json.dump(jobs, f, indent=2, default=str)
        print(f"Successfully saved {len(jobs)} jobs")
    except Exception as e:
        print(f"Error saving jobs: {e}")

@app.route('/', methods=['GET', 'POST'])
@login_required
def index():
    print(f"DEBUG: index() called with method={request.method}")
    global processing_complete, current_output_dir
    if request.method == 'POST':
        print(f"POST request received at {datetime.now()}")
        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 job ID first
        job_start_time = datetime.now()
        job_id = job_start_time.strftime('%Y%m%d_%H%M%S')
        
        # Create dirs
        input_dir = os.path.join(os.path.dirname(__file__), 'uploads', job_id)
        os.makedirs(input_dir, exist_ok=True)
        output_dir = tempfile.mkdtemp()
        current_output_dir = output_dir
        
        for file in files:
            if file and file.filename:
                try:
                    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_path = os.path.join(input_dir, filename)
                        file.save(file_path)
                        print(f"Saved PDF: {filename} to {file_path}")
                    else:
                        # Assume folder upload
                        file_path = os.path.join(input_dir, filename)
                        file.save(file_path)
                        print(f"Saved file from folder: {filename} to {file_path}")
                except Exception as e:
                    print(f"Error saving file {file.filename}: {e}")
                    flash(f'Error saving file {file.filename}: {e}')
                    return redirect(request.url)
        
        # Create job immediately
        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)
        print(f"Job {job_id} created and saved at {job_start_time}")
        
        # Start processing in background
        processing_complete = False
        thread = threading.Thread(target=process_pdfs_async, args=(input_dir, output_dir, job_id))
        thread.start()
        print(f"Thread started for job {job_id}")
        
        # 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('/test')
def test():
    try:
        test_file = os.path.join(os.path.dirname(__file__), 'uploads', 'test.txt')
        with open(test_file, 'w') as f:
            f.write('test')
        return f'Test file created at {test_file}'
    except Exception as e:
        return f'Error: {e}'
    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():
    jobs = load_jobs()
    print(f"Loaded {len(jobs)} jobs: {list(jobs.keys())}")
    # Sort jobs by timestamp, newest first
    sorted_jobs = sorted(jobs.values(), key=lambda x: x.get('timestamp', datetime.min), reverse=True)
    print(f"Sorted jobs count: {len(sorted_jobs)}")
    if sorted_jobs:
        print(f"First job: {sorted_jobs[0].get('id')} status: {sorted_jobs[0].get('status')}")
    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, job_id=None):
    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 (only if PDF support is available)
    if PDF_SUPPORT:
        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)
    else:
        msg = "Step 3/4: PDF merging skipped (PyMuPDF not available)\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        print(msg.strip())
        time.sleep(0.1)
    
    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):
    logger.info("=" * 60)
    logger.info(f"PROCESS_PDFS_ASYNC CALLED")
    logger.info(f"  Job ID: {job_id}")
    logger.info(f"  PDF Dir: {pdfdir}")
    logger.info(f"  Output Folder: {outputfolder}")
    logger.info(f"  PDF Dir Exists: {os.path.exists(pdfdir)}")
    if os.path.exists(pdfdir):
        logger.info(f"  PDF Dir Writable: {os.access(pdfdir, os.W_OK)}")
        logger.info(f"  PDF Dir Contents: {os.listdir(pdfdir)}")
    logger.info(f"  Output Dir Exists: {os.path.exists(outputfolder)}")
    if os.path.exists(outputfolder):
        logger.info(f"  Output Dir Writable: {os.access(outputfolder, os.W_OK)}")
    logger.info(f"  PDF Support Available: {PDF_SUPPORT}")
    logger.info(f"  Timestamp: {datetime.now()}")
    logger.info("=" * 60)
    
    global processing_complete
    
    if not PDF_SUPPORT:
        error_msg = "PyMuPDF (fitz) is not installed - PDF processing cannot continue"
        logger.error(f"ERROR: {error_msg}")
        jobs = load_jobs()
        if job_id in jobs:
            jobs[job_id]['status'] = 'failed'
            jobs[job_id]['error'] = error_msg
            jobs[job_id]['message'] = error_msg
            save_jobs(jobs)
        return
    
    try:
        jobs = load_jobs()
        logger.info(f"Jobs loaded: {len(jobs)} total, job {job_id} exists: {job_id in jobs}")
        
        if job_id not in jobs:
            logger.error(f"ERROR: Job {job_id} not found 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")
        logger.info("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, exist_ok=True)
        os.makedirs(missing_dir, exist_ok=True)
        logger.info(f"Created output directories: {addresses_dir}, {missing_dir}")
        
        # Scan for PDF files
        pdf_files = []
        logger.info(f"Scanning directory: {pdfdir}")
        for root, dirs, files in os.walk(pdfdir):
            logger.info(f"  Checking directory: {root} with {len(files)} files")
            for f in files:
                if f.endswith('.pdf'):
                    pdf_path = os.path.join(root, f)
                    pdf_files.append(pdf_path)
                    logger.info(f"    Found PDF: {f}")
        
        logger.info(f"Total PDFs found: {len(pdf_files)}")
        
        if len(pdf_files) == 0:
            error_msg = f"No PDF files found in {pdfdir}"
            logger.error(f"ERROR: {error_msg}")
            jobs = load_jobs()
            jobs[job_id]['status'] = 'failed'
            jobs[job_id]['error'] = error_msg
            jobs[job_id]['message'] = error_msg
            save_jobs(jobs)
            return
        
        msg = f"Step 1/4: Found {len(pdf_files)} PDF files\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        logger.info(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)
        logger.info("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
        logger.info(f"Processing {len(pdf_files)} files sequentially (avoiding multiprocessing issues)")
        results = []
        for i, pdf_file in enumerate(pdf_files, 1):
            logger.info(f"Processing file {i}/{len(pdf_files)}: {os.path.basename(pdf_file)}")
            try:
                result = process_single_pdf(pdf_file)
                results.append(result)
                logger.info(f"  Result type: {result['type']}")
            except Exception as e:
                logger.error(f"  ERROR processing {pdf_file}: {e}")
                results.append({
                    'type': 'error',
                    'filename': os.path.basename(pdf_file),
                    'error': str(e),
                    'filepath': pdf_file
                })
            
            # Update progress periodically
            if i % 10 == 0 or i == len(pdf_files):
                progress_percent = 20 + int((i / len(pdf_files)) * 60)  # 20-80%
                jobs = load_jobs()
                jobs[job_id]['message'] = f"Processing file {i}/{len(pdf_files)}"
                jobs[job_id]['progress'] = progress_percent
                save_jobs(jobs)
                logger.info(f"  Progress: {progress_percent}%")
        
        # Process results sequentially
        logger.info("Processing results and sorting files...")
        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):
            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']))
                if 'pages' in result:
                    pages = result['pages']
                    if pages not in page_groups:
                        page_groups[pages] = []
                    page_groups[pages].append(result['filepath'])
            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']))
            elif result['type'] == 'error':
                logger.error(f"  Error with {result['filename']}: {result.get('error', 'Unknown error')}")
        
        logger.info(f"Results sorted. Found {len(page_groups)} page count groups for merging")
        
        msg = "Step 2/4: Sorting complete. Starting merging...\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        logger.info(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 (only if PDF support is available)
        if PDF_SUPPORT:
            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)
            logger.info(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)
                logger.info(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)
                    logger.info(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)
                    logger.error(msg.strip())
                time.sleep(0.05)
        else:
            msg = "Step 3/4: PDF merging skipped (PyMuPDF not available)\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            print(msg.strip())
            time.sleep(0.1)
        
        msg = f"Processing complete! Processed {len(pdf_files)} files total\n"
        with open(status_file, 'a') as status_f:
            status_f.write(msg)
        logger.info(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)
        logger.info(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)
            logger.info(f"Job {job_id} completed successfully")
            
            msg = f"ZIP created: {zip_filename}\n"
            with open(status_file, 'a') as status_f:
                status_f.write(msg)
            logger.info(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)
            logger.error(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:
        logger.error(f"Processing failed for job {job_id}: {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
    logger.info(f"Processing finished for job {job_id} at {datetime.now()}")

@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."""
    if not PDF_SUPPORT:
        return {'type': 'error', 'filename': os.path.basename(filepath), 'error': 'PDF processing not available (PyMuPDF not installed)', 'filepath': filepath}
    
    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=100*1024*1024*1024)  # 100GB
    except ImportError:
        print("Waitress not available, using Flask dev server...")
        app.run(host='0.0.0.0', port=5000, debug=True)