from flask import Flask, render_template, request, jsonify, send_file, Response
from parser import parse_source_to_graph
from dataset_gen import create_dataset_entry, get_dataset_stats, upload_to_hub, build_dataset_entry, OUTPUT_FILE
import os
import json
from blueprints.summarize import summarize_bp

app = Flask(__name__)
app.register_blueprint(summarize_bp, url_prefix='/summarize/')

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

@app.route('/parse', methods=['POST'])
def parse():
    code = request.json.get('code', '')
    return jsonify(parse_source_to_graph(code))

# --- New API Endpoint ---
@app.route('/api/process', methods=['POST'])
def api_process():
    """
    Direct API for external tools.
    Input: JSON { "code": "...", "format": "json" | "jsonl" }
    Output: The dataset entry object or a raw JSONL string.
    """
    data = request.get_json()
    if not data or 'code' not in data:
        return jsonify({"error": "No code provided"}), 400
        
    code = data['code']
    fmt = data.get('format', 'json') # Default to JSON object
    
    # Generate the entry without saving to disk
    entry = build_dataset_entry(code)
    
    if "error" in entry:
        return jsonify(entry), 400
    
    if fmt == 'jsonl':
        # Return as a string suitable for appending to a file
        return Response(json.dumps(entry) + '\n', mimetype='text/plain')
    
    # Return as standard JSON object
    return jsonify(entry)


@app.route('/dataset/add', methods=['POST'])
def add_entry():
    code = request.json.get('code', '')
    return jsonify(create_dataset_entry(code))

@app.route('/dataset/list', methods=['GET'])
def list_entries():
    return jsonify(get_dataset_stats())

@app.route('/dataset/download', methods=['GET'])
def download_dataset():
    if os.path.exists(OUTPUT_FILE):
        return send_file(OUTPUT_FILE, as_attachment=True)
    return "No dataset found", 404

@app.route('/dataset/upload_hf', methods=['POST'])
def upload_hf():
    data = request.json
    return jsonify(upload_to_hub(data.get('token'), data.get('repo_id')))

if __name__ == '__main__':
    app.run(host="0.0.0.0", port=7860, debug=True)