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

LOGO_B64 = "iVBORw0KGgoAAAANSUhEUgAAB4AAAALOCAYAAAE/zz7rAAEY70lEQVR42uzdd7wcVf3/8aEXCSIgKBDTG+RuuQkBUkhC6J3QO6GDGAiBkEJCEkok9AC hBBBEBb9fwIICUkVAxfpTREAQpX1VikAsoJT7+5zduWSz2TLlnDNnZl5/PB+Ee+/uzpyZz3zmPW29crnTy7A/2v7MwulfBqzJcvF21aCAQQGnrGgpYFDAFDAFDAo4yeKlgEEBp7h4KWBQwBQwBQwKmAIGKGAKGMhLAb9HAYMCdtd32xTwPRQwKOD07j7fSgGDAk5vAZOBQQFTwBQwKGAKGKCAKWAgCwV8V4Di/YgCBgWc3u57IgUMCpjdZwoYuS7gcwwUFgUMCtiRTpkWZ1DAyGMHzlIRU8CggFPsNQoYeSvgz2asiLsoYOTtKHSWinflFeZv2wmsVKCAE/JcrHkbNZYVCrk4D5yWDnt44HkauwMrEyjgtB1xpvOCK7FWNFj8j8UC/Yu4RHwm9LSOHs9KBArYUCf/q9HP325nViBQwAYL+A2dn9dx7k2sMKCALRbwq9qKd+4SVhZQwJYL+EUtNydQvKCAHS3gUskrD+zvlQcN8EpjtvOKk6Z4hdnXeIUzLmLlAAXs4i40Cx0UcDoKeCkFDAo4Y0/jYKGDAqaAAQqYAgYo4OAFPHgACx0UcFoLmAUOCjhbBdwVUffrP6z7+cQWg93sPRr5MODnx/2cqK+J8hnKvyLOU9RlF/S9Jmt6n0bT1UvDfCgvUMDpKOCg07BeiNeHXYE8S6+pd5vmAo46/3EK2eTnr/B+FLA7BbxRhOm4IcQ8mV7hbBaJF2el11U8DdxvqYC7KGD3CrhL0wqlYwU2XcBxu3mU8XpU03gviDhNugu4K38FPO0S3XlQVwHr7JhBV6KeBla2oF2p1fu+oSm3t5q+Rr97XsMeio4CXjPEe++Xr9NIncPSUsCe5o7WrafhbtHoNYdoOGAVdTkFLeAPQn7W1AB/89m6/980ZgE3HP+sF/BuK7xuUH+TBbyfWKWJZu97lOaDSEGKrEtTAXsxd9WXOlLA7T7r4wgdP+o0rhlmvv5zYQD+5sq4CjZNUoGnBVy5TRxwCfMeP1H48GruLvQ6xs41fa9ut/dTAGb/mbCAf1cLeAbY65QYQvYVjf7tCNHoacYKOA48YgCjvpI2BR34BsiFFScAq7//Y9ijNnKBs+3BnnP1ShgCrhZAY+IsALsrCEDvx+xgKOuZLquzLrdQgGbuPIsztj1poDdLWBdR6FXttANehk+R93sPYLMz+oaI8RnDVy8srPGseEodAYLWDm+wXusHWGFqP3dJA27gHEjwrUWrt82fbFL9+++bbGAG73P1RSwGwVs8kos3cWgI+PHPQag48KYqPM1RNdVVFyJla0CjrIwOw0XcJji2lbjSmriNJKJa6F1XE7KtdAZKmBPw0pusoBtdHLT54GjjvedmnP0F7gbKZsFrFzTYuH92uL9wLpvMzyxxXyNbfSeMZdT/d9spPE65rBjc3LMi2xm8Uid8JdVmjdlgVfuGBr99VMXZu9JE2pMapZT8dgzefoGT+SIdFnlcjrmXq93JZWbKz6ZjmHDg7/2rMu8zv59lq3gBx0f+HXq7zt7bxZ5usvlcuO9makXai/c0EU86yr/GzfKredBxroyDgP7BZ+20y6ovKa0677h5mnGFUHvggu4wV42zmr9oYDbFni/ZZrtevuScJKPeeaFf9WFmbTz2pyt9Ryhd5sBT94+SJWXxcTbZ7L1a+YiTAdcb6SprzVyPbzeOy05YtdfTVO9+8LhUARqbt4l329Trnuu7BW3Eh3zLuh+Rffzbm2biN0fnXs1Vf2NJuXfr1bjIXaI+tYNs8nnd1w3rpRwFG/9Ozsq73C9Mva/133Aq4U7rBg713Zyi6QLts73HRtuXXrFSeMmo1JedTYSK8L1nUvMLqedPbr5XfCReGW78xFldd19u0V7DWjxnmlsduHX49kQ/RJ4Q4dGvr1FPDyxta9/pYI0xNl+u+IMb86pum2EH+r7d7r8pBB7d4/yrhHHZcwr4uyfjV/XeMOG2i+KODljYjxPt1/+7MY059EARubjo7ZK0aI4sEnNHvf/4k49lEK+FVNBazs0uY1B5ucLwq4cQFH6apRCn6VGJ+j62/r/+ZoXe8dsoCjLrMkCjhut/coYHsdOOxgh3mN50AB7xdxzEwU8GYa1wOXCvh0U/NFATcu4Ili35C7Oy+ELLAuDRsLE39bu2dgu4DHWSrglzXuQgeZvhspYHd3obtiLFhPw66V7l3oqBnNxC60+v+5Bqd/nRhjFWec1P9fyS60Gwexmi2gn4bIwC4UsPY9gZAFvJrFg1iexs/6isXPooDb6BDP1Ng8wGvU363V4GfPBHhdlAIO8t5R/jbsmLV975AFXD8dN0WcjjDz/KOQ8x3ns7o/5x6dn0UBw5iIBQwu5KCAKWAKmAIGBUwBU8AUMAVMAVPAFDAoYFDAFDAFDAqYAqaAKWBQwBQwBUwBU8CggCnguAW4Vd3//zhgoTZzbcTLFUEB566AjxTvxSw4W1j4FHAuCzhIcXR30evFX1r83fsGC/RscW+DW7+iFHr3v5eywlDAaSngjVus1C/7v3dhBj4T8u83qNnA1Bf2oiY/f6jBGOzPykMBu1TAs8XqOR+MJ2oKeA//333FpU0Kex//v92d/gBWKAqYo9Dpobr/qPKyZ0o9XlPkGzI+FDAFnO5TXP+siSMUMAVMAYMCBgUMCpgCBihgCti0Xax+3vARfGct+H5gDY5J5MorChgUsJEbHChgUMCO+06T4n2EAgYFTPelgEEBJ3BvMAUMCpgCpoBBAZvwbwoYFHB2H61DAYMCpoApYFDAFDBAAVPAAAVMAYMCTtQSChgUMA93p4BBAVPAFDAoYAoYoIApYFDAFDAFDAqYAgYoYAoYoIANFfDmg1ngoIDTWsCFGVc0GoCuiJq9vtlAB/27INMV9rXt/v6kkK85KMJntJuvX4dccZcGWEZ/0DBdyh4hp03nsm77Pvkp4FPPS0MBd8VdoA3+7sgIK5DOv1deijlPUQq420WGxzvu2HTFmY78FPBpF7hewF1xF2aElahLUwFP17iS6i5gZRuD421iPNuZlrsCLp440+UC7tIwLVFWoi5NBd8VsSs1+t1CAwVcP333axzvJAq4K4dHoYeZyDAmCzhqd41b6KYLuP73R2vowGNCfm7Q6f+Hho1hlz+9cdahNZpNR75OI01Z4GIB686gQYrsgwjFuFfd36xsYJ7D1nAOfq/v8JQ5u8ytIFr+bdZLuCoRRxnQFdpIkzBxd3ItFqJ3o6xsjX6m7grts0C7tvk529o2KhHHZu461tX7gq4WsQXmCrgKDnKZgF3aSxE9bM/1v1sVsRx+lBzAf9Jw3GLuAUcZ1wDT0suC7hNJ85aAXsaCljH9Cd1GslEAU+kgJP+atFiiQKO32FMXLhgooBtj3e31+p+9y4F7HYB614BdL5fqyuVWr3fzjELXPdpl7AFbOJiC0/Txi+3GfitAMX7QkoLWPdR6LgroY5zpfu0mK/rNBTwwBDL5AiN423ivH2YaemR1gJ+PEABn5+SAjZ9HrjR7zzDBTwv5nzFOYhV/16vaB7vlxKIUJ1ZOw98WYACHpjyAraxRQ/btXUcNFIe9n93qZed6DEhx+BwQ8Vn/T3SWsDTAxTweikqYJPXQuss4Cjv8SkL10KPKei7Yi1q4W0QYmzejDkdC/JQwFumrIBN3Y0UtYC315gVpwWcp3UsXUrZ7SmLezq69nIycTfSOQEK+BspLOAk7gcOeuqpYOCCh5bTWTzmDK9w9tW6CljnNd5dCRRwpu4HvlHLV4MGuzbaqNJBx0efDvUg+gw+ZaI8oB+PQgrinOtSW8CP6vp2hlDFM3m+vsLd+9DlpqNj7pLAry0eNy3SCl4uFoNcStpUx/wbG4/joP56CneLzWN9G+Qnt47K+DSdh3lLgl5Su5zOXpvI92QNCzdPZy/2OuZcq22dWe69pHjVPKS1gJ8p2/h+JNnNLg/sVyXfk1RZ6A2KuDx4YOAvRStvNbL5Sjpj0YqvOeuy8F93Ku8TeB5L5eX+v7TzPpXPXGE65t3Q+n3iPHNMCinsPBaPrdmIFQor/G3x6NNXLIC5SxrMf+mTf3dK52+4zKRwg85zedCAqohfTVuU9av2q38KUy9s+V5pLeB3kvyO4MIpcyvZrGHhLlfEy7bY5W1GB3rvzn69/CK8Ivy0DR+hbR5L2+9WnY4zLqrcSx34td2vC9p1x4w3+z3O51xf/ZyB/QO/prOvvwymB1wGas9h1LjQ09bZr/eywpV1Ksr85fdSymDvG+W1Jj4nyvy9EXF6tEyL2t1uWbyNN1DvRBzH+p/Nivi6MPN+R4h149ma160dY75epoDDv2+Q130jZgEr62gu4CdjFKSWjUmrnKppHput6CdZKGCTG99m0/cjCjjcwJnoYM0+J8hrF4f4jDAFfHXI+Q00zSEKuPu9Vqn52W9jdCqTBTzF//ftIQpxq5qf/TXGfFHAhjpwV8hCrP27Pg4UcJQxM1HAutYDkwV8RIhp1j1fFHDA951Y8+9NQiykBw3uipkqYhYQACIAJgAGAAAAABAAEwADAACAABgAAQAAQABMAAwAAAACYAAAAACRW0isL/YS24k9xD5BuJ/wmkG4xzDlBAAgACYABgAAAAEwAAAAgIRYWgwXk4JwX+FulAkAgACYABgAAAAEwAAAAACSOYO4h1hKzEV5AAAIgAmAAQAAQAAMAAAAAAAAEAATAAMAAIAAGAAAAAAAACAABgAAAAiAAQAAALhmTrGA/S/lAQAgACYABgAAQIz+P1GZ7DqsPrIUAAAAAElFTkSuQmCC"

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
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'] = '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('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()

@app.context_processor
def inject_logo():
    return {'logo_b64': LOGO_B64}

# 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('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."""
    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('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:
                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
        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('/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():
    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):
    print(f"process_pdfs_async called with job_id={job_id}, pdfdir={pdfdir}, outputfolder={outputfolder}")
    global processing_complete
    print(f"Processing started for job {job_id} at {datetime.now()}")
    try:
        jobs = load_jobs()
        if job_id not in jobs:
            print(f"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")
        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
        print(f"Processing {len(pdf_files)} files sequentially")
        results = []
        for pdf_file in pdf_files:
            result = process_single_pdf(pdf_file)
            results.append(result)
        
        # 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 (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)
            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)
            print(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)
            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:
        print(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
    print(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=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)