from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import os
from parser import parse_python_code

app = Flask(__name__)
UPLOAD_FOLDER = 'uploads'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Ensure upload folder exists
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/parse_code', methods=['POST'])
def parse_code():
    code = None
    if 'file' in request.files and request.files['file'].filename:
        file = request.files['file']
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)
        with open(file_path, 'r') as f:
            code = f.read()
        os.remove(file_path)  # Clean up
    elif 'code' in request.form:
        code = request.form['code']

    if not code:
        return jsonify({'error': 'No code or file provided'}), 400

    parts, _ = parse_python_code(code)
    
    # Extract nodes (functions and variables) and connections
    nodes = []
    connections = []
    node_id_map = {}  # Maps parser node_id to canvas node index
    y_offset = 50  # For vertical layout

    for part in parts:
        category = part['category']
        node_id = part['node_id']
        source = part['source'].strip()

        if category in ['function', 'assigned_variable', 'input_variable', 'returned_variable']:
            node_data = {
                'id': len(nodes),
                'type': category,
                'label': node_id,
                'source': source,
                'x': 50,
                'y': y_offset,
                'inputs': [],
                'outputs': []
            }

            if category == 'function':
                # Extract function name and parameters
                if 'def ' in source:
                    func_name = source.split('def ')[1].split('(')[0]
                    node_data['label'] = func_name
                    params = source.split('(')[1].split(')')[0].split(',')
                    params = [p.strip() for p in params if p.strip()]
                    node_data['inputs'] = params
                    # Assume function can have outputs (return statements)
                    node_data['outputs'] = ['return']
            elif category == 'input_variable':
                # Input variable (function parameter)
                var_name = source.strip().rstrip(',')
                node_data['label'] = var_name
                node_data['outputs'] = [var_name]
            elif category == 'assigned_variable':
                # Assigned variable
                var_name = source.split('=')[0].strip()
                node_data['label'] = var_name
                node_data['inputs'] = ['value']
                node_data['outputs'] = [var_name]
            elif category == 'returned_variable':
                # Returned variable
                var_name = source.split('return ')[1].strip() if 'return ' in source else node_id
                node_data['label'] = var_name
                node_data['inputs'] = [var_name]

            nodes.append(node_data)
            node_id_map[node_id] = node_data['id']
            y_offset += 100

    # Create connections based on parent_path and variable usage
    for part in parts:
        if part['category'] in ['input_variable', 'assigned_variable', 'returned_variable']:
            parent_path = part['parent_path']
            if 'Function' in parent_path:
                parent_node_id = parent_path.split(' -> ')[0]
                if parent_node_id in node_id_map:
                    from_id = node_id_map[part['node_id']]
                    to_id = node_id_map[parent_node_id]
                    if part['category'] == 'input_variable':
                        connections.append({'from': from_id, 'to': to_id})
                    elif part['category'] == 'returned_variable':
                        connections.append({'from': to_id, 'to': from_id})

    return jsonify({'nodes': nodes, 'connections': connections})

@app.route('/save_nodes', methods=['POST'])
def save_nodes():
    data = request.get_json()
    nodes = data.get('nodes', [])
    connections = data.get('connections', [])
    return jsonify({'status': 'success', 'nodes': nodes, 'connections': connections})

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