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 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', 'import']:
            node_data = {
                'id': len(nodes),
                'type': category,
                'label': node_id,
                'source': source,
                'x': 50,
                'y': y_offset,
                'inputs': [],
                'outputs': []
            }

            if category == 'function':
                # Function: inputs are parameters, output is return
                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
                    node_data['outputs'] = ['return']
            elif category == 'input_variable':
                # Input variable: output to parent function
                var_name = source.strip().rstrip(',')
                node_data['label'] = var_name
                node_data['outputs'] = [var_name]
            elif category == 'assigned_variable':
                # Assigned variable: input is value, output is 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: input from function
                var_name = source.split('return ')[1].strip() if 'return ' in source else node_id
                node_data['label'] = var_name
                node_data['inputs'] = [var_name]
            elif category == 'import':
                # Import: no inputs/outputs, just a node
                import_name = source.split('import ')[1].split()[0] if 'import ' in source else node_id
                node_data['label'] = import_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:
        category = part['category']
        if category in ['input_variable', 'returned_variable', 'assigned_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 category == 'input_variable':
                        connections.append({'from': from_id, 'to': to_id})
                    elif category == 'returned_variable':
                        connections.append({'from': to_id, 'to': from_id})
                    elif category == 'assigned_variable':
                        # Connect assigned variables to parent function (e.g., local variables)
                        connections.append({'from': from_id, 'to': to_id})
        elif category == 'import':
            # Connect imports to functions that might use them (simplified: connect to all functions)
            for other_part in parts:
                if other_part['category'] == 'function' and other_part['node_id'] in node_id_map:
                    connections.append({
                        'from': node_id_map[part['node_id']],
                        'to': node_id_map[other_part['node_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(host="0.0.0.0", port=7860, debug=True)