# app.py
from flask import Flask, request, render_template
from parser import parse_python_code
import os

app = Flask(__name__)

def reconstruct_code(parts):
    """Reconstruct the original code from parsed parts."""
    sorted_parts = sorted(parts, key=lambda p: p['location'][0])
    return ''.join(part['source'] for part in sorted_parts)

@app.route('/', methods=['GET', 'POST'])
def index():
    if request.method == 'POST':
        parts = None
        filename = 'unnamed.py'
        code_input = None

        # Handle file upload
        if 'file' in request.files and request.files['file'].filename:
            file = request.files['file']
            if not file.filename.endswith('.py'):
                return 'Invalid file type. Please upload a Python file.', 400
            filename = file.filename
            file_path = os.path.join('uploads', filename)
            file.save(file_path)
            with open(file_path, 'r') as f:
                code_input = f.read()
            parts = parse_python_code(code_input)

        # Handle pasted code
        elif 'code' in request.form and request.form['code'].strip():
            code_input = request.form['code']
            filename = request.form.get('filename', 'unnamed.py') or 'unnamed.py'
            if not filename.endswith('.py'):
                filename += '.py'
            parts = parse_python_code(code_input)

        if parts:
            indexed_parts = [{'index': i + 1, **part} for i, part in enumerate(parts)]
            reconstructed_code = reconstruct_code(indexed_parts)
            # Return partial HTML for htmx to update the results section
            return render_template(
                'results_partial.html',
                parts=indexed_parts,
                filename=filename,
                reconstructed_code=reconstructed_code,
                code_input=code_input
            )
        return 'No file or code provided', 400

    # Initial page load
    return render_template('index.html', parts=None, filename=None, reconstructed_code=None, code_input=None)

if __name__ == '__main__':
    if not os.path.exists('uploads'):
        os.makedirs('uploads')
    app.run(host="0.0.0.0",port=7860)