import os
from flask import Flask, request, render_template, jsonify
from werkzeug.utils import secure_filename
import logging
from pathlib import Path
from utils.feature_extractor import ShapeFeatureExtractor
from utils.classifier import ShapeClassifier
import cv2  # Add this import

app = Flask(__name__)

# Configuration
app.config.update(
    UPLOAD_FOLDER=Path('static/uploads'),
    CLASSIFIED_FOLDER=Path('static/classified_results'),
    ANNOTATIONS_FOLDER=Path('static/annotations'),
    MAX_CONTENT_LENGTH=16 * 1024 * 1024  # 16MB max file size
)

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create required directories
for folder in [app.config['UPLOAD_FOLDER'], 
               app.config['CLASSIFIED_FOLDER'], 
               app.config['ANNOTATIONS_FOLDER']]:
    folder.mkdir(parents=True, exist_ok=True)

ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

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

@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        if 'file' not in request.files:
            return render_template('upload.html', error="No file part")
            
        files = request.files.getlist('file')
        results = {}
        
        # Initialize feature extractor and classifier
        feature_extractor = ShapeFeatureExtractor(
            use_advanced_features=request.form.get('advanced_features', 'true') == 'true'
        )
        classifier = ShapeClassifier(
            app.config['CLASSIFIED_FOLDER'],
            similarity_threshold=float(request.form.get('similarity', '85')) / 100
        )
        
        for file in files:
            if file and allowed_file(file.filename):
                try:
                    filename = secure_filename(file.filename)
                    filepath = app.config['UPLOAD_FOLDER'] / filename
                    file.save(filepath)
                    
                    # Extract features and classify
                    features, annotated_img = feature_extractor.extract_features(filepath)
                    if features:
                        # Save annotated image
                        annotated_path = app.config['ANNOTATIONS_FOLDER'] / filename
                        cv2.imwrite(str(annotated_path), annotated_img)
                        
                        # Classify image
                        class_folder, similarity = classifier.classify_image(filepath, features)
                        
                        # Move file to classified folder
                        classified_path = app.config['CLASSIFIED_FOLDER'] / class_folder
                        classified_path.mkdir(exist_ok=True)
                        filepath.rename(classified_path / filename)
                        
                        results[filename] = {
                            'features': features,
                            'annotated': str(annotated_path),
                            'class': class_folder,
                            'similarity': similarity
                        }
                        
                except Exception as e:
                    logger.error(f"Error processing {filename}: {str(e)}")
                    continue
                    
        return render_template('result.html', results=results)
    
    return render_template('upload.html')

@app.route('/dashboard')
def dashboard():
    classifier = ShapeClassifier(app.config['CLASSIFIED_FOLDER'])
    folders = {}
    total_images = 0
    
    # Iterate through all folders in the classified_results directory
    for folder in app.config['CLASSIFIED_FOLDER'].iterdir():
        if folder.is_dir():
            # Get all image files in the folder
            images = list(folder.glob('*.{jpg,jpeg,png}'))
            folders[folder.name] = {
                'path': str(folder),
                'count': len(images),
                'images': [img.name for img in images[:5]]  # Preview first 5 images
            }
            total_images += len(images)
            
    return render_template('dashboard.html', 
                         folders=folders, 
                         total_images=total_images)

@app.route('/reclassify/<folder>', methods=['POST'])
def reclassify_folder(folder):
    classifier = ShapeClassifier(app.config['CLASSIFIED_FOLDER'])
    success = classifier.reclassify_folder(folder)
    return jsonify({'success': success})

@app.route('/merge', methods=['POST'])
def merge_folders():
    data = request.get_json()
    classifier = ShapeClassifier(app.config['CLASSIFIED_FOLDER'])
    success = classifier.merge_classes(data['source'], data['target'])
    return jsonify({'success': success})

if __name__ == '__main__':
    app.run(debug=True)