from flask import Flask, render_template, request, jsonify, send_file, Response, stream_with_context
from werkzeug.utils import secure_filename
import os
from pathlib import Path
import shutil
import io
import json
import logging
from backend.pipeline import classify

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

app = Flask(__name__)

# Configuration
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
UPLOAD_FOLDER_SINGLE = 'static/uploads/single'
UPLOAD_FOLDER_MULTIPLE = 'static/uploads/multiple'

app.config['UPLOAD_FOLDER_SINGLE'] = UPLOAD_FOLDER_SINGLE
app.config['UPLOAD_FOLDER_MULTIPLE'] = UPLOAD_FOLDER_MULTIPLE
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024
app.config['MAX_CONTENT_LENGTH_REPORT'] = 500 * 1024 * 1024

# Ensure upload directories exist
os.makedirs(UPLOAD_FOLDER_SINGLE, exist_ok=True)
os.makedirs(UPLOAD_FOLDER_MULTIPLE, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def clear_uploads(folder):
    """Helper function to clear upload directories."""
    if os.path.exists(folder):
        for filename in os.listdir(folder):
            file_path = os.path.join(folder, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                logger.error(f'Failed to delete {file_path}. Reason: {e}')

@app.route('/', defaults={'path': ''})
@app.route('/<path:path>')
def index(path):
    if path.startswith('static/'):
        return send_file(path)
    return send_file('static/react/index.html')

@app.route('/upload_single', methods=['POST'])
def upload_single():
    if 'file' not in request.files:
        return jsonify({'error': 'No file uploaded'}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No file selected'}), 400
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER_SINGLE'], filename)
        file.save(filepath)
        return jsonify({'filename': filename})
    return jsonify({'error': 'Invalid file type'}), 400

@app.route('/classify_single', methods=['POST'])
def classify_single():
    data = request.get_json()
    filename = data.get('filename')
    if not filename:
        return jsonify({'error': 'No filename provided'}), 400
    filepath = os.path.join(app.config['UPLOAD_FOLDER_SINGLE'], filename)
    if not os.path.exists(filepath):
        return jsonify({'error': 'File not found'}), 404
    
    try:
        classification_result, detailed_results, failure_labels = classify(filepath)
        return jsonify({
            'classification': classification_result,
            'detailed_results': detailed_results
        })
    except Exception as e:
        logger.error(f"Error in classify_single: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/upload_multiple', methods=['POST'])
def upload_multiple():
    logger.info("=== UPLOAD_MULTIPLE CALLED ===")
    if 'file' not in request.files:
        logger.warning("No 'file' in request.files")
        return jsonify({'error': 'No file uploaded'}), 400
    
    files = request.files.getlist('file')
    logger.info(f"Received {len(files)} files in request")
    if not files:
        return jsonify({'error': 'No files selected'}), 400

    try:
        # Ensure temp directory exists (DON'T wipe it)
        temp_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
        os.makedirs(temp_dir, exist_ok=True)
        
        saved_count = 0
        filename_map = {}
        
        # Load existing map if present
        map_path = os.path.join(temp_dir, 'filename_map.json')
        if os.path.exists(map_path):
            try:
                with open(map_path, 'r') as f:
                    filename_map = json.load(f)
                logger.info(f"Loaded existing filename map with {len(filename_map)} entries")
            except Exception as e:
                logger.error(f"Error loading existing filename map: {e}")
        
        for file in files:
            if file and allowed_file(file.filename):
                original_filename = file.filename
                filename = secure_filename(file.filename)
                filepath = os.path.join(temp_dir, filename)
                file.save(filepath)
                filename_map[filename] = original_filename
                saved_count += 1
                logger.info(f"Saved: '{original_filename}' -> '{filename}'")
        
        # Save updated filename map
        with open(map_path, 'w') as f:
            json.dump(filename_map, f)
        
        # Count actual files in directory
        actual_files = [f for f in os.listdir(temp_dir) if allowed_file(f)]
        logger.info(f"Total files in temp directory after upload: {len(actual_files)}")
        logger.info(f"Files: {actual_files}")
        
        return jsonify({
            'message': f'Successfully uploaded {saved_count} files',
            'count': saved_count,
            'status': 'Ready'
        })
        
    except Exception as e:
        logger.error(f"Error in upload_multiple: {e}", exc_info=True)
        return jsonify({'error': str(e)}), 500

@app.route('/classify_multiple', methods=['POST'])
def classify_multiple():
    logger.info("=== CLASSIFY_MULTIPLE CALLED ===")
    def generate():
        temp_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
        if not os.path.exists(temp_dir):
            logger.warning("Temp directory does not exist")
            yield json.dumps({'error': 'No files to classify'}) + '\n'
            return

        # Load filename map
        filename_map = {}
        map_path = os.path.join(temp_dir, 'filename_map.json')
        if os.path.exists(map_path):
            try:
                with open(map_path, 'r') as f:
                    filename_map = json.load(f)
                logger.info(f"Loaded filename map: {filename_map}")
            except Exception as e:
                logger.error(f"Error loading filename map: {e}")

        files = [f for f in os.listdir(temp_dir) if allowed_file(f)]
        logger.info(f"Processing {len(files)} files from temp directory")
        logger.info(f"Files to process: {files}")
        
        for filename in files:
            filepath = os.path.join(temp_dir, filename)
            logger.info(f"Classifying: {filename}")
            try:
                classification_result, _, failure_labels = classify(filepath)
                logger.info(f"Result for {filename}: {classification_result} with labels: {failure_labels}")
                
                # Move file
                dest_dir = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], classification_result.lower())
                os.makedirs(dest_dir, exist_ok=True)
                dest_path = os.path.join(dest_dir, filename)
                shutil.move(filepath, dest_path)
                
                # Get original filename if available
                original_filename = filename_map.get(filename, filename)
                logger.info(f"Sending result for original filename: {original_filename}")
                
                result = {
                    'filename': original_filename,
                    'status': 'pass' if classification_result == 'Pass' else 'fail',
                    'labels': failure_labels,
                    'score': 0
                }
                yield json.dumps(result) + '\n'
                
            except Exception as e:
                logger.error(f"Error processing {filename}: {e}", exc_info=True)
                # Use original filename for error reporting
                original_filename = filename_map.get(filename, filename)
                yield json.dumps({'filename': original_filename, 'status': 'error', 'error': str(e)}) + '\n'

    return Response(stream_with_context(generate()), mimetype='application/x-ndjson')

@app.route('/clear_uploads', methods=['POST'])
def clear_uploads_route():
    logger.info("=== CLEAR_UPLOADS CALLED ===")
    try:
        clear_uploads(app.config['UPLOAD_FOLDER_SINGLE'])
        clear_uploads(app.config['UPLOAD_FOLDER_MULTIPLE'])
        return jsonify({'success': True})
    except Exception as e:
        logger.error(f"Error in clear_uploads_route: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/use_sample', methods=['POST'])
def use_sample():
    try:
        data = request.get_json()
        filename = data.get('filename')
        destination = data.get('destination') # 'single' or 'multiple'
        
        if not filename or not destination:
            return jsonify({'error': 'Missing filename or destination'}), 400
            
        # Validate filename (security)
        if not allowed_file(filename):
            return jsonify({'error': 'Invalid filename'}), 400
            
        # Source path
        src_path = os.path.join('static', 'samples', filename)
        if not os.path.exists(src_path):
            return jsonify({'error': 'Sample not found'}), 404
            
        # Destination path
        if destination == 'single':
            dest_folder = app.config['UPLOAD_FOLDER_SINGLE']
        elif destination == 'multiple':
            dest_folder = os.path.join(app.config['UPLOAD_FOLDER_MULTIPLE'], 'temp')
            os.makedirs(dest_folder, exist_ok=True)
        else:
            return jsonify({'error': 'Invalid destination'}), 400
            
        dest_path = os.path.join(dest_folder, filename)
        
        # Copy file
        shutil.copy2(src_path, dest_path)
        
        # For multiple, we need to update the filename map
        if destination == 'multiple':
            map_path = os.path.join(dest_folder, 'filename_map.json')
            filename_map = {}
            if os.path.exists(map_path):
                try:
                    with open(map_path, 'r') as f:
                        filename_map = json.load(f)
                except:
                    pass
            
            filename_map[filename] = filename # Map to itself for samples
            
            with open(map_path, 'w') as f:
                json.dump(filename_map, f)
                
        return jsonify({'success': True, 'filename': filename})
        
    except Exception as e:
        logger.error(f"Error in use_sample: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/samples', methods=['GET'])
def get_samples():
    try:
        samples_dir = os.path.join('static', 'samples')
        if not os.path.exists(samples_dir):
            return jsonify([])
            
        files = [f for f in os.listdir(samples_dir) if allowed_file(f)]
        # Sort files for consistent order
        files.sort()
        
        samples = []
        for i, filename in enumerate(files):
            samples.append({
                'id': i + 1,
                'url': f'/static/samples/{filename}',
                'filename': filename
            })
            
        return jsonify(samples)
    except Exception as e:
        logger.error(f"Error in get_samples: {e}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    logger.info("SERVER STARTING ON PORT 7860")
    # Disable reloader to prevent loading models twice (saves memory)
    app.run(debug=False, use_reloader=False, host='0.0.0.0', port=7860)