app.py

import os
import json
import threading
import logging
import pymongo
import gradio as gr
from flask import Flask, request, jsonify
from flask_cors import CORS
from geopy.distance import geodesic
import colorlog

# --- Logging Setup ---
handler = colorlog.StreamHandler()
handler.setFormatter(colorlog.ColoredFormatter(
    "%(log_color)s[%(levelname)s] %(asctime)s - %(message)s",
    datefmt='%Y-%m-%d %H:%M:%S',
    log_colors={
        'DEBUG': 'cyan',
        'INFO': 'green',
        'WARNING': 'yellow',
        'ERROR': 'red',
        'CRITICAL': 'bold_red'
    }
))
logger = colorlog.getLogger('flask_app')
logger.addHandler(handler)
logger.setLevel(logging.DEBUG)

# --- MongoDB Setup ---
mongo_uri = os.getenv(
    "MONGO_URI",
    "mongodb+srv://township_chatbot:Daniel%[email protected]/db?retryWrites=true&w=majority&appName=Cluster0"
)
mongo_client = pymongo.MongoClient(mongo_uri)
db = mongo_client["township_db"]
businesses_collection = db["businesses"]

# --- Flask Setup ---
app = Flask(__name__)
app.config["DEBUG"] = True
CORS(app)

# --- MCP Logic ---
def register_business(params):
    required = ["name", "description", "address", "phone", "latitude", "longitude"]
    if not all(k in params for k in required):
        logger.warning(f"Missing required fields: {params}")
        return {"error": f"Missing required fields. Got: {params}"}
    businesses_collection.insert_one(params)
    logger.info(f"Registered new business: {params['name']}")
    return {"status": "success", "message": "Business registered"}

def get_nearby_businesses(params):
    try:
        lat = float(params.get("latitude"))
        lon = float(params.get("longitude"))
        radius_km = float(params.get("radius_km", 10))
    except (TypeError, ValueError):
        logger.error("Invalid or missing latitude/longitude")
        return {"error": "Invalid or missing latitude/longitude"}

    all_businesses = list(businesses_collection.find({}))
    nearby = []
    for b in all_businesses:
        try:
            b_lat = float(b["latitude"])
            b_lon = float(b["longitude"])
            dist = geodesic((lat, lon), (b_lat, b_lon)).km
            if dist <= radius_km:
                b["distance_km"] = round(dist, 2)
                b["_id"] = str(b["_id"])
                nearby.append(b)
        except Exception as e:
            logger.warning(f"Skipping business with error: {e}")
            continue

    logger.info(f"Found {len(nearby)} nearby businesses within {radius_km} km")
    return {"businesses": nearby}

def mcp_server(json_input):
    try:
        request_obj = json.loads(json_input)
        method = request_obj.get("method")
        params = request_obj.get("params", {})
        logger.debug(f"Received method: {method} with params: {params}")

        if method == "registerBusiness":
            result = register_business(params)
        elif method == "getNearbyBusinesses":
            result = get_nearby_businesses(params)
        else:
            result = {"error": f"Unknown method: {method}"}

        response = {"jsonrpc": "2.0", "id": request_obj.get("id"), "result": result}
    except Exception as e:
        logger.exception("Error in MCP server logic")
        response = {"jsonrpc": "2.0", "error": str(e)}
    return json.dumps(response)

# --- Flask Route for JSON-RPC POST ---
@app.route("/run/predict", methods=["POST"])
def run_predict():
    try:
        incoming_data = request.get_json()
        logger.info(f"Incoming JSON-RPC: {json.dumps(incoming_data, indent=2)}")

        if not incoming_data or 'data' not in incoming_data:
            logger.error("Missing 'data' field")
            return jsonify({"jsonrpc": "2.0", "error": "Missing 'data' field"}), 400

        raw_jsonrpc = incoming_data['data'][0]
        rpc_obj = json.loads(raw_jsonrpc)
        logger.debug(f"Parsed JSON-RPC Object: {json.dumps(rpc_obj, indent=2)}")

        method = rpc_obj.get("method")
        params = rpc_obj.get("params", {})

        if method == "registerBusiness":
            result = register_business(params)
        elif method == "getNearbyBusinesses":
            result = get_nearby_businesses(params)
        else:
            logger.warning(f"Unknown method: {method}")
            return jsonify({"jsonrpc": "2.0", "error": f"Unknown method: {method}"}), 400

        response = {"jsonrpc": "2.0", "id": rpc_obj.get("id"), "result": result}
        logger.info(f"Response: {json.dumps(response, indent=2)}")
        return jsonify(response)

    except json.JSONDecodeError as e:
        logger.exception("JSON decode error")
        return jsonify({"jsonrpc": "2.0", "error": f"JSON decode error: {str(e)}"}), 400

    except Exception as e:
        logger.exception("Unhandled server error")
        return jsonify({"jsonrpc": "2.0", "error": f"Server error: {str(e)}"}), 500

# --- Gradio UI Launch ---
flag_reasons = [
    "Incorrect response",
    "Output is confusing",
    "Incomplete answer",
    "Other"
]

def launch_gradio():
    try:
        iface = gr.Interface(
            fn=mcp_server,
            inputs=[gr.Textbox(label="JSON-RPC Input", lines=20)],
            outputs=[gr.Textbox(label="Chatbot Response", lines=20)],
            title="Township Business Chatbot UI",
            description="Send JSON-RPC commands to the backend",
            flagging_dir="flagged",
            flagging_options=flag_reasons
        )
        iface.launch(server_name="0.0.0.0", server_port=7860)
    except OSError as e:
        logger.error(f"⚠️ Gradio failed to launch: {str(e)}")

# --- Main Entrypoint ---
if __name__ == "__main__":
    threading.Thread(target=launch_gradio, daemon=True).start()
    logger.info("Starting Flask server...")
    app.run(host="0.0.0.0", port=5000, debug=True)