app.py

from flask import Flask, render_template, request, send_file, jsonify
from werkzeug.utils import secure_filename
from pathlib import Path
import io
import os
import shutil
from dotenv import load_dotenv
import google.genai as genai
from markitdown import MarkItDown

# Load environment variables from .env file
load_dotenv()

app = Flask(__name__)

# Configuration
UPLOAD_FOLDER = Path(__file__).parent / 'tmp'
IMAGE_EXTENSIONS = {'jpg', 'jpeg', 'png', 'gif', 'bmp', 'webp'}
ALLOWED_EXTENSIONS = {
    'pdf', 'docx', 'doc', 'pptx', 'xlsx', 'xls', 'csv', 'json', 'xml',
    'html', 'htm', 'txt', 'epub', 'msg', 'zip', 'jpg', 'jpeg', 'png', 'gif',
    'bmp', 'webp', 'mp3', 'wav', 'm4a', 'flac'
}

app.config['UPLOAD_FOLDER'] = str(UPLOAD_FOLDER)
app.config['MAX_CONTENT_LENGTH'] = 100 * 1024 * 1024  # 100MB max file size

# Initialize Gemini 2.0 Flash for image OCR only
gemini_client = None
google_api_key = os.getenv('GOOGLE_API_KEY', '')

if google_api_key:
    try:
        gemini_client = genai.Client(api_key=google_api_key)
        print("✓ Gemini 2.0 Flash initialized for image OCR")
    except Exception as e:
        print(f"Warning: Failed to initialize Gemini: {e}")
else:
    print("Note: No GOOGLE_API_KEY found. Image OCR will not be available.")

# Initialize MarkItDown for non-image files
md = MarkItDown()


def cleanup_temp_folder():
    """Delete temporary folder and all its contents"""
    if UPLOAD_FOLDER.exists():
        shutil.rmtree(UPLOAD_FOLDER)


def allowed_file(filename):
    """Check if file extension is allowed"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def is_image_file(filename):
    """Check if file is an image"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in IMAGE_EXTENSIONS


def get_supported_formats():
    """Return list of supported file formats"""
    return sorted(list(ALLOWED_EXTENSIONS))


def ocr_image_with_gemini(filepath):
    """Extract text from image using Gemini 2.0 Flash"""
    if not gemini_client:
        return None
    
    try:
        from PIL import Image
        
        # Open and prepare the image
        image = Image.open(filepath)
        
        try:
            # Call Gemini API with vision capability
            response = gemini_client.models.generate_content(
                model="gemini-2.0-flash",
                contents=[
                    image,
                    "Extract all text from this image. Be thorough and preserve formatting where possible."
                ]
            )
            return response.text
        finally:
            # Always close the image to release the file lock
            image.close()
    except Exception as e:
        error_msg = str(e)
        # Check for quota exceeded errors (429)
        if "429" in error_msg or "quota" in error_msg.lower():
            print(f"Gemini quota exceeded. OCR not available for this image.")
            return None
        else:
            print(f"Error during Gemini OCR: {e}")
            return None


@app.route('/')
def index():
    """Render the main page"""
    formats = get_supported_formats()
    return render_template('index.html', formats=formats)


@app.route('/api/formats', methods=['GET'])
def formats():
    """Return supported file formats as JSON"""
    return jsonify({'formats': get_supported_formats()})


@app.route('/api/convert', methods=['POST'])
def convert():
    """Convert uploaded file to markdown"""
    try:
        # Check if file is in request
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': f'File type not supported. Allowed types: {", ".join(ALLOWED_EXTENSIONS)}'}), 400
        
        # Create temp folder, save uploaded file
        UPLOAD_FOLDER.mkdir(exist_ok=True)
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            markdown_content = ""
            
            # Handle images with Gemini OCR
            if is_image_file(filename):
                markdown_content = ocr_image_with_gemini(filepath)
                if not markdown_content:
                    markdown_content = f"# {Path(filename).stem}\n\n"
                    markdown_content += f"**File:** {filename}\n"
                    markdown_content += f"**Type:** {Path(filename).suffix.lower()}\n\n"
                    markdown_content += "**Note:** Image OCR service temporarily unavailable (quota exceeded). Please try again later or enable billing on your Google Cloud project.\n"
            else:
                # Use MarkItDown for other file types
                result = md.convert(filepath)
                markdown_content = result.text_content

                if not markdown_content or markdown_content.strip() == '':
                    markdown_content = f"# {Path(filename).stem}\n\n"
                    markdown_content += f"**File:** {filename}\n"
                    markdown_content += f"**Type:** {Path(filename).suffix.lower()}\n\n"
                    
                    # Add metadata if available
                    if hasattr(result, 'metadata') and result.metadata:
                        markdown_content += "## Metadata\n\n"
                        for key, value in result.metadata.items():
                            markdown_content += f"- **{key}:** {value}\n"
                    else:
                        markdown_content += "*No extractable content found in this file.*\n"
            
            # Return markdown as downloadable file
            output = io.BytesIO()
            output.write(markdown_content.encode('utf-8'))
            output.seek(0)
            
            # Generate output filename with same name but .md extension
            output_filename = Path(filename).stem + '.md'
            
            return send_file(
                output,
                mimetype='text/markdown',
                as_attachment=True,
                download_name=output_filename
            )
        
        finally:
            # Clean up temp folder completely
            cleanup_temp_folder()
    
    except Exception as e:
        return jsonify({'error': f'Conversion error: {str(e)}'}), 500


@app.route('/api/preview', methods=['POST'])
def preview():
    """Get preview of markdown conversion"""
    try:
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
        
        file = request.files['file']
        
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        if not allowed_file(file.filename):
            return jsonify({'error': f'File type not supported'}), 400
        
        # Create temp folder, save uploaded file
        UPLOAD_FOLDER.mkdir(exist_ok=True)
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            markdown_content = ""
            
            # Handle images with Gemini OCR
            if is_image_file(filename):
                markdown_content = ocr_image_with_gemini(filepath)
                if not markdown_content:
                    markdown_content = "**Note:** Image OCR service temporarily unavailable (quota exceeded). Please try again later or enable billing on your Google Cloud project."
            else:
                # Use MarkItDown for other file types
                result = md.convert(filepath)
                markdown_content = result.text_content
                if not markdown_content or markdown_content.strip() == '':
                    markdown_content = "*No extractable content found in this file.*"
            
            # Return first 2000 characters as preview
            preview_text = markdown_content[:2000]
            if len(markdown_content) > 2000:
                preview_text += '\n\n... (content truncated for preview)'
            
            return jsonify({'preview': preview_text})
        
        finally:
            # Clean up temp folder completely
            cleanup_temp_folder()
    
    except Exception as e:
        return jsonify({'error': f'Preview error: {str(e)}'}), 500


if __name__ == '__main__':
    app.run(debug=True, host='127.0.0.1', port=5000)