app.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import os
import subprocess # For running shell commands
import platform   # For detecting the operating system
import time       # For small delays
import zipfile    # For handling ZIP files
import shutil     # For file operations
import uuid

from prompt import get_new_json_single_edit, get_new_json_general, summarize_changes, rl_prompt

app = Flask(__name__)
CORS(app)  # Enable CORS for all routes

# Create uploads directory if it doesn't exist
UPLOADS_DIR = "uploads"
if not os.path.exists(UPLOADS_DIR):
    os.makedirs(UPLOADS_DIR)

@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', 'http://localhost:3000')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
    response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
    response.headers.add('Access-Control-Allow-Credentials', 'true')
    return response

# --- Upload Agent ZIP File Endpoint ---
@app.route('/upload_agent', methods=['POST'])
def upload_agent():
    if 'agent_zip' not in request.files:
        return jsonify({"error": "No ZIP file part in the request"}), 400

    zip_file = request.files['agent_zip']

    if zip_file.filename == '':
        return jsonify({"error": "No selected ZIP file"}), 400

    if not zip_file.filename.endswith('.zip'):
        return jsonify({"error": "File must be a ZIP archive"}), 400

    try:
        # Create a unique directory name for this upload in root directory
        upload_id = str(uuid.uuid4())
        upload_path = os.path.join(os.getcwd(), upload_id)

        # Delete existing folder if it exists
        if os.path.exists(upload_path):
            shutil.rmtree(upload_path)

        # Create the directory
        os.makedirs(upload_path, exist_ok=True)

        # Save the ZIP file temporarily
        zip_path = os.path.join(upload_path, zip_file.filename)
        zip_file.save(zip_path)

        # Extract the ZIP file
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(upload_path)

        # Remove the ZIP file after extraction
        os.remove(zip_path)

        # Look for agent.json or similar files in the extracted content
        agent_files = []
        for root, dirs, files in os.walk(upload_path):
            for file in files:
                if file.endswith('.json') and ('agent' in file.lower() or 'config' in file.lower()):
                    agent_files.append(os.path.join(root, file))

        return jsonify({
            "success": True,
            "message": f"Agent uploaded successfully: {zip_file.filename}",
            "upload_id": upload_id,
            "agent_files": agent_files,
            "extracted_path": upload_path
        }), 200

    except zipfile.BadZipFile:
        return jsonify({"error": "Invalid ZIP file format"}), 400
    except Exception as e:
        return jsonify({"error": f"An error occurred during upload: {str(e)}"}), 500

# --- Configuration for ADK Web Server ---
ADK_WEB_PORT = 8000  # Default port for 'adk web'. Adjust if yours is different.

def call_single_edit_agent(input_json_data, prompt, param='root_agent'):
    """
    This function simulates your custom Gemini agent processing.
    In a real application, you would replace this with your
    actual Gemini agent's interaction, which might involve:
    - Calling a Gemini API
    - Using a local ML model
    - Applying business rules based on the prompt

    Args:
        input_json_data (dict): The parsed JSON data from the input file.
        prompt (str): The prompt provided by the user.

    Returns:
        tuple: (updated_json_data (dict), context_of_changes (str))
    """
    updated_json_data = input_json_data.copy()
    context_of_changes = get_new_json_single_edit(input_json_data, param, prompt)
    # TODO: need to take in param to edit

    # Example: Modify JSON based on a simple prompt
    if "add_timestamp" in prompt.lower():
        import datetime
        updated_json_data["last_updated"] = datetime.datetime.now().isoformat()
        context_of_changes += "- Added 'last_updated' timestamp.\n"

    if "change_status_to_processed" in prompt.lower():
        if "status" in updated_json_data:
            updated_json_data["status"] = "processed"
            context_of_changes += "- Changed 'status' to 'processed'.\n"
        else:
            updated_json_data["status"] = "newly_processed"
            context_of_changes += "- Added 'status' as 'newly_processed'.\n"

    if "add_notes" in prompt.lower() and "notes_content" in prompt.lower():
        # Extract notes content from prompt (a more robust solution would use regex or a more structured prompt)
        try:
            notes_start = prompt.find("notes_content:") + len("notes_content:")
            notes_end = prompt.find("'", notes_start) # Assuming notes_content ends with a single quote for simplicity
            if notes_end == -1: # if no closing quote found, take till end
                notes_content = prompt[notes_start:].strip()
            else:
                notes_content = prompt[notes_start:notes_end].strip()
            updated_json_data["notes"] = notes_content
            context_of_changes += f"- Added notes: '{notes_content}'.\n"
        except Exception as e:
            context_of_changes += f"- Failed to add notes due to error: {e}.\n"


    if not context_of_changes:
        context_of_changes = "No specific changes requested or applied based on the prompt."

    return updated_json_data, context_of_changes

def call_general_agent(input_json_data, prompt):
    updated_json_str, changelog = get_new_json_general(input_json_data, prompt)
    updated_json_data = json.loads(updated_json_str)
    context_of_changes = changelog
    return updated_json_data, context_of_changes


# --- Helper Function to find and kill processes on a port ---
def _kill_process_on_port(port):
    print(f"Attempting to kill processes on port {port}...")
    if platform.system() == "Windows":
        print(f"  Platform: Windows")
        try:
            print(f"  Running netstat command...")
            # subprocess.check_output returns bytes by default, so decode it.
            output = subprocess.check_output(f"netstat -ano | findstr :{port}", shell=True).decode('utf-8')
            print(f"  netstat output received. Parsing PIDs...")
            pids = [line.strip().split()[-1] for line in output.splitlines() if line.strip()]
            for pid in set(pids):
                print(f"  Attempting to kill PID: {pid}")
                # taskkill also returns bytes, decode stderr if an error occurs.
                subprocess.run(f"taskkill /F /PID {pid}", shell=True, check=True, capture_output=True)
            print(f"  Successfully killed processes on port {port} (Windows).")
            return True
        except subprocess.CalledProcessError as e:
            # FIX: Check if e.stderr is None before decoding
            error_output = "No stderr output captured." # Default message if stderr is None
            if e.stderr is not None:
                error_output = e.stderr.decode('utf-8', errors='ignore')
            print(f"  Error running netstat/taskkill (Windows): {error_output.strip()}")
            return False
        except Exception as e:
            print(f"  Unexpected error in _kill_process_on_port (Windows): {e}")
            return False
    else: # macOS and Linux
        print(f"  Platform: macOS/Linux")
        try:
            print(f"  Running lsof command: sudo lsof -ti :{port}")
            lsof_cmd = f"sudo lsof -ti :{port}"
            # text=True captures stdout as a string. stderr will still be captured for CalledProcessError.
            pids_output = subprocess.check_output(lsof_cmd, shell=True, text=True).strip()
            print(f"  lsof output: '{pids_output}'")
            if not pids_output:
                print(f"  No process found on port {port}.")
                return True # No process to kill is a success

            pids = pids_output.splitlines()
            for pid in pids:
                print(f"  Attempting to kill PID: {pid} with 'sudo kill -9 {pid}'")
                # capture_output=True ensures stderr is captured as bytes.
                subprocess.run(f"sudo kill -9 {pid}", shell=True, check=True, capture_output=True)
            print(f"  Successfully killed processes on port {port} (macOS/Linux).")
            return True
        except subprocess.CalledProcessError as e:
            # FIX: Check if e.stderr is None before decoding
            error_output = "No stderr output captured." # Default message if stderr is None
            if e.stderr is not None:
                error_output = e.stderr.decode('utf-8', errors='ignore')
            print(f"  Error running lsof/kill (macOS/Linux): {error_output.strip()}")
            print("  This might be due to lack of sudo password or the command not being found.")
            return False
        except Exception as e:
            print(f"  Unexpected error in _kill_process_on_port (macOS/Linux): {e}")
            return False


# --- Helper Function to run ADK web in background ---
def _start_adk_web_in_background(adk_port):
    print(f"Starting 'adk web' in background on port {adk_port}...")

    cmd = f'adk web --port {adk_port}'

    creationflags = 0
    is_windows = platform.system() == "Windows"
    if is_windows:
        creationflags = subprocess.DETACHED_PROCESS

    try:
        # Use Popen to run in background and detach.
        # On Windows, close_fds cannot be True when redirecting standard handles.
        process = subprocess.Popen(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                                   stdin=subprocess.DEVNULL, close_fds=not is_windows,
                                   creationflags=creationflags)
        print(f"  'adk web' process started with PID: {process.pid}")
        return True
    except Exception as e:
        print(f"  Failed to start 'adk web': {e}")
        return False

def _run_adk_query(query):
    # Build the command
    cmd = f"printf '{query}' | adk run hotels_com_api_agent"

    try:
        # Run the command and capture output
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        output = result.stdout + (result.stderr if result.stderr else '')

        return jsonify({
            # 'output_file': output_filename,
            'output': output
        }), 200
    except Exception as e:
        return jsonify({'error': f'Failed to run command: {str(e)}'}), 500


# --- API Endpoint 1: Process JSON Single Edit ---
@app.route('/process_json/single_edit/<session_id>/<node_id>', methods=['POST'])
def process_json_single_edit(session_id, node_id):
    if 'json_file' not in request.files:
        return jsonify({"error": "No JSON file part in the request"}), 400

    json_file = request.files['json_file']
    prompt = request.form.get('prompt')
    param = request.form.get('param', 'root_agent')  # Default to 'root_agent' if not provided

    if json_file.filename == '':
        return jsonify({"error": "No selected JSON file"}), 400

    if not prompt:
        return jsonify({"error": "No prompt provided"}), 400

    if json_file:
        try:
            # Read the JSON file content
            json_data_str = json_file.read().decode('utf-8')
            input_json_data = json.loads(json_data_str)

            # --- Call your custom Gemini agent logic for single edit ---
            updated_json, context_of_changes = call_single_edit_agent(input_json_data, prompt, param)

            # Create session history directory if it doesn't exist
            session_dir = os.path.join('history', session_id)
            if not os.path.exists(session_dir):
                os.makedirs(session_dir)

            # Save the updated JSON to a file as <node_id>.json
            file_path = os.path.join(session_dir, f"{node_id}.json")
            with open(file_path, 'w') as f:
                json.dump(updated_json, f, indent=2)

            # Summarize the change for node naming
            node_name = summarize_changes(prompt, context_of_changes)

            return jsonify({
                "updated_json": json.dumps(updated_json, indent=2),
                "context_of_changes": context_of_changes,
                "node_name": node_name
            }), 200

        except json.JSONDecodeError:
            return jsonify({"error": "Invalid JSON file format"}), 400
        except Exception as e:
            return jsonify({"error": f"An error occurred: {str(e)}"}), 500
    else:
        return jsonify({"error": "An unexpected error occurred with the file upload"}), 500

# --- API Endpoint 2: Process JSON General ---
@app.route('/process_json/general/<session_id>/<node_id>', methods=['POST'])
def process_json_general(session_id, node_id):
    if 'json_file' not in request.files:
        return jsonify({"error": "No JSON file part in the request"}), 400

    json_file = request.files['json_file']
    prompt = request.form.get('prompt')

    if json_file.filename == '':
        return jsonify({"error": "No selected JSON file"}), 400

    if not prompt:
        return jsonify({"error": "No prompt provided"}), 400

    if json_file:
        try:
            # Read the JSON file content
            json_data_str = json_file.read().decode('utf-8')
            input_json_data = json.loads(json_data_str)

            # --- Call your custom Gemini agent logic for general ---
            updated_json, context_of_changes = call_general_agent(input_json_data, prompt)

            # Create session history directory if it doesn't exist
            session_dir = os.path.join('history', session_id)
            if not os.path.exists(session_dir):
                os.makedirs(session_dir)

            # Save the updated JSON to a file as <node_id>.json
            file_path = os.path.join(session_dir, f"{node_id}.json")
            with open(file_path, 'w') as f:
                json.dump(updated_json, f, indent=2)

            # Summarize the change for node naming
            node_name = summarize_changes(prompt, context_of_changes)

            return jsonify({
                "updated_json": json.dumps(updated_json, indent=2),
                "context_of_changes": context_of_changes,
                "node_name": node_name
            }), 200

        except json.JSONDecodeError:
            return jsonify({"error": "Invalid JSON file format"}), 400
        except Exception as e:
            return jsonify({"error": f"An error occurred: {str(e)}"}), 500
    else:
        return jsonify({"error": "An unexpected error occurred with the file upload"}), 500

# --- API Endpoint 3: Get Node JSON in Session ---
@app.route('/history/<session_id>/<node_id>', methods=['GET'])
def get_node_json(session_id, node_id):
    session_dir = os.path.join('history', session_id)
    file_path = os.path.join(session_dir, f"{node_id}.json")

    if not os.path.exists(file_path):
        return jsonify({"error": "Node JSON not found for the given session and node ID"}), 404

    try:
        with open(file_path, 'r') as f:
            data = json.load(f)
        return jsonify(data)
    except Exception as e:
        return jsonify({"error": f"An error occurred while reading the node JSON file: {str(e)}"}), 500

# --- API Endpoint 4: Update Agent JSON File ---
@app.route('/update_agent_json', methods=['POST'])
def update_agent_json():
    try:
        data = request.get_json()
        if not data or 'json_data' not in data:
            return jsonify({"error": "No JSON data provided"}), 400

        # Path to the agent.json file
        agent_json_path = os.path.join("hotels_com_api_agent", "agent.json")

        # Write the new JSON data to the file
        with open(agent_json_path, 'w') as f:
            json.dump(data['json_data'], f, indent=2)

        return jsonify({
            "status": "success",
            "message": "Agent JSON file updated successfully"
        }), 200

    except Exception as e:
        return jsonify({"error": f"An error occurred while updating the agent JSON file: {str(e)}"}), 500

# --- API Endpoint 5: Retrigger ADK Web Server ---
@app.route('/retrigger_adk_web', methods=['POST'])
def retrigger_adk_web():
    print(f"Received request to retrigger ADK web server on port {ADK_WEB_PORT}.")

    # 1. Attempt to kill existing ADK web server process
    kill_success = _kill_process_on_port(ADK_WEB_PORT)
    if not kill_success:
        print("  Warning: Failed to kill existing ADK web process, attempting restart anyway.")

    # Give a tiny moment for the port to free up if a process was just killed
    time.sleep(1)

    # 2. Start new ADK web server process
    start_success = _start_adk_web_in_background(ADK_WEB_PORT)

    if start_success:
        return jsonify({
            "status": "success",
            "message": f"Attempted to restart ADK web server on port {ADK_WEB_PORT}. Check ADK logs for status."
        }), 200
    else:
        return jsonify({
            "status": "error",
            "message": "Failed to start ADK web server. Check server logs for details."
        }), 500

# --- API Endpoint 5: Create a new session ---
@app.route('/new_session', methods=['POST'])
def new_session():
    """
    Create a new session directory under history/ with a random id and return the id.
    """
    session_id = str(uuid.uuid4())
    session_dir = os.path.join('history', session_id)
    os.makedirs(session_dir, exist_ok=True)
    return jsonify({"session_id": session_id}), 200

# --- API Endpoint 6: RL to loop ADK Query and Save Output ---
@app.route("/rl", methods=["POST"])
def rl():
    data = request.get_json()
    query = data.get('query')
    if not query:
        return jsonify({'error': 'No query provided'}), 400

    agent_json_path = os.path.join('hotels_com_api_agent', 'agent.json')

    # 1. Read current agent.json
    with open(agent_json_path, 'r') as f:
        agent_json = json.load(f)

    # 2. Run ADK query and get output
    adk_response, status_code = _run_adk_query(query)
    if status_code != 200:
        return adk_response, status_code
    adk_output = adk_response.get_json().get('output', '')

    # 3. Generate RL prompt
    rl_instruction = rl_prompt(query, adk_output)


    # 4. Use LLM to update agent.json
    updated_json_str, _ = get_new_json_general(agent_json, rl_instruction)
    try:
        updated_json = json.loads(updated_json_str)
    except Exception as e:
        return jsonify({'error': f'Failed to parse updated agent JSON: {e}', 'raw': updated_json_str}), 500

    # 5. Overwrite agent.json
    with open(agent_json_path, 'w') as f:
        json.dump(updated_json, f, indent=2)

    return jsonify({'status': 'success', 'message': 'RL step completed.'}), 200

if __name__ == '__main__':
    # You can change the host and port for your Flask app as needed
    # Ensure this port is DIFFERENT from ADK_WEB_PORT (e.g., 5000 for Flask, 8000 for ADK)
    app.run(debug=True, host='0.0.0.0', port=5000)