from flask import Flask, request, jsonify, abort
import requests
import os
import xml.etree.ElementTree as ET
import time

app = Flask(__name__)

# Optional token for access control (not production secure)
HEALTH_TOKEN = os.environ.get('HEALTH_TOKEN')

@app.route('/healthcheck')
def healthcheck():
    token = request.args.get('token') or request.headers.get('X-Health-Token')
    if HEALTH_TOKEN and token != HEALTH_TOKEN:
        return jsonify({'error': 'Unauthorized'}), 401

    # New behavior: prefer service id lookup
    service_id = request.args.get('id')
    host = request.args.get('host')
    port = request.args.get('port')
    proto = request.args.get('proto', 'http')

    if service_id:
        # Load services.xml and find service by id
        try:
            tree = ET.parse('/usr/share/nginx/html/services.xml')
            root = tree.getroot()
        except Exception as e:
            return jsonify({'error': 'unable to parse services.xml', 'detail': str(e)}), 500

        # Find service element by id
        svc = None
        for s in root.findall('.//service'):
            sid = s.get('id') or s.get('name')
            if sid and sid == service_id:
                svc = s
                break
        if svc is None:
            return jsonify({'error': 'service id not found'}), 404

        # Determine host/port/proto preference: local-ip > tailscale-ip root > host attr > fallback to provided host
        local_ip = svc.get('local-ip')
        root_tailscale = root.get('tailscale-ip')
        host_attr = svc.get('host')
        proto_attr = svc.get('proto') or proto
        port_attr = svc.get('port') or port
        health_path = svc.get('health-path', '')  # Custom health check path (e.g., /api/health, /ping)

        # prefer local-ip if present
        target_host = local_ip or root_tailscale or (host_attr.split(':')[0] if host_attr and ':' in host_attr else None) or host
        target_port = port_attr or (host_attr.split(':')[1] if host_attr and ':' in host_attr else None)
        target_proto = proto_attr
        if not target_host or not target_port:
            return jsonify({'error': 'service missing host or port mapping'}), 400

        host = target_host
        port = target_port
        proto = target_proto

    if not host or not port:
        return jsonify({'error': 'missing parameters, expected host and port or id'}), 400

    # Build URL with optional custom health path
    health_path = health_path if 'health_path' in locals() else ''
    url = f"{proto}://{host}:{port}{health_path}"

    try:
        # Use HEAD to do a lightweight check
        r = requests.head(url, timeout=4, allow_redirects=True, verify=False)
        # Any 2xx/3xx is considered OK
        ok = 200 <= r.status_code < 400
        return jsonify({'ok': ok, 'status_code': r.status_code}), (200 if ok else 502)
    except requests.RequestException as e:
        return jsonify({'ok': False, 'error': str(e)}), 502

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8081)