from flask import Flask, render_template, request, jsonify, Response, stream_with_context
from flask_cors import CORS
from datetime import datetime
import os
import json
import time
import threading

app = Flask(__name__)
CORS(app)

# ── In-memory state ───────────────────────────────────────────────────────────
# The Cloudflare Worker calls /api/update-urls whenever the Electron app
# re-registers its cloudflared tunnel URL (every 1 s for the first 30 s,
# then every 10 s as a slow heartbeat).
current_urls = {
    'base_url': '',
    'last_updated': '',
    'portals': []
}

# SSE subscriber queues — each connected browser gets its own queue.
# When /api/update-urls is called, we push the new data to all queues.
_sse_lock = threading.Lock()
_sse_clients: list = []   # list of threading.Queue


def _broadcast_sse(data: dict):
    """Push a JSON payload to every connected SSE client."""
    msg = f"data: {json.dumps(data)}\n\n"
    with _sse_lock:
        dead = []
        for q in _sse_clients:
            try:
                q.put_nowait(msg)
            except Exception:
                dead.append(q)
        for q in dead:
            _sse_clients.remove(q)


# ── Window definitions ────────────────────────────────────────────────────────
# Paths MUST match the React Router routes in frontend/src/App.tsx
WINDOW_DEFINITIONS = [
    {
        'window': 1,
        'name': 'AEGIS Portal',
        'path': '/',
        'icon': 'home',
        'color': 'blue',
        'description': 'Main AEGIS portal — login and system overview',
    },
    {
        'window': 2,
        'name': 'Tech Analysis',
        'path': '/window/2',
        'icon': 'cpu',
        'color': 'blue',
        'description': 'TEC model — technology threat scores for any year',
    },
    {
        'window': 3,
        'name': 'Conductor Synthesis',
        'path': '/window/3',
        'icon': 'git-branch',
        'color': 'purple',
        'description': 'Groq synthesis of tech threats into strategic insights',
    },
    {
        'window': 4,
        'name': 'Economic Analysis',
        'path': '/window/4',
        'icon': 'trending-up',
        'color': 'yellow',
        'description': '12 economic impact indicators from conductor results',
    },
    {
        'window': 5,
        'name': 'War Analysis',
        'path': '/window/5',
        'icon': 'shield',
        'color': 'red',
        'description': 'Dual AI + live conflict data — war risk predictions',
    },
    {
        'window': 6,
        'name': 'Climate Analysis',
        'path': '/window/6',
        'icon': 'cloud',
        'color': 'cyan',
        'description': 'World Bank CCKP — climate-conflict nexus assessment',
    },
    {
        'window': 7,
        'name': 'Disease Prediction',
        'path': '/window/7',
        'icon': 'activity',
        'color': 'orange',
        'description': 'AI biosecurity risk — top 3 disease outbreak predictions',
    },
    {
        'window': 8,
        'name': 'Drug Discovery',
        'path': '/window/8',
        'icon': 'zap',
        'color': 'green',
        'description': '4-phase pipeline — SMILES, DiffDock, Vina, manufacturing',
    },
    {
        'window': 9,
        'name': 'Visual Disease',
        'path': '/window/9',
        'icon': 'image',
        'color': 'pink',
        'description': 'AI-generated pathogen morphology visualizations',
    },
    {
        'window': 10,
        'name': 'Visual Kinetic',
        'path': '/window/10',
        'icon': 'eye',
        'color': 'indigo',
        'description': '3D PK/PD simulation with live cardiac monitoring',
    },
    {
        'window': 11,
        'name': 'AR Portal',
        'path': '/window/11',
        'icon': 'scan',
        'color': 'teal',
        'description': 'Augmented Reality — scan QR to view all windows remotely',
    },
]


def build_portals(base_url: str) -> list:
    """Build portal list from window definitions + a given base URL."""
    return [
        {
            **w,
            'url': f"{base_url.rstrip('/')}{w['path']}",
            'qr_url': f"{base_url.rstrip('/')}{w['path']}",
        }
        for w in WINDOW_DEFINITIONS
    ]


# ── Routes ────────────────────────────────────────────────────────────────────

@app.route('/')
def index():
    """AR Portal Dashboard — shows QR grid for all windows."""
    return render_template(
        'ar_dashboard.html',
        base_url=current_urls.get('base_url', ''),
        last_updated=current_urls.get('last_updated', ''),
        portals=current_urls.get('portals', []),
    )


@app.route('/ping', methods=['GET'])
def ping():
    """UptimeRobot keep-alive."""
    return jsonify({'status': 'ok', 'message': 'AR Space is alive'}), 200


@app.route('/health')
def health():
    return jsonify({'status': 'healthy', 'timestamp': datetime.now().isoformat()})


@app.route('/api/status')
def status():
    return jsonify({
        'status': 'online',
        'base_url': current_urls.get('base_url', ''),
        'last_updated': current_urls.get('last_updated', ''),
        'portals_count': len(current_urls.get('portals', [])),
        'total_windows': len(WINDOW_DEFINITIONS),
    })


@app.route('/api/update-urls', methods=['POST'])
def update_urls():
    """
    Called by the Cloudflare Worker every ~1 s (burst) then every 10 s.
    Updates in-memory state and pushes to all SSE subscribers instantly.
    """
    try:
        data = request.get_json()
        if not data or 'base_url' not in data:
            return jsonify({'error': 'base_url is required'}), 400

        base_url  = data['base_url'].rstrip('/')
        portals   = build_portals(base_url)
        timestamp = datetime.now().isoformat()

        current_urls.update({
            'base_url':     base_url,
            'last_updated': timestamp,
            'portals':      portals,
        })

        # Push to all connected SSE clients immediately
        _broadcast_sse({
            'base_url':  base_url,
            'timestamp': timestamp,
            'portals':   portals,
        })

        return jsonify({
            'success':       True,
            'message':       'URLs updated successfully',
            'base_url':      base_url,
            'timestamp':     timestamp,
            'portals_count': len(portals),
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/stream')
def sse_stream():
    """
    Server-Sent Events endpoint.
    The dashboard connects here and receives live updates whenever
    /api/update-urls is called (i.e. every 1 s during the 30 s burst
    and every 10 s after that).
    """
    import queue as _queue

    q = _queue.Queue(maxsize=50)
    with _sse_lock:
        _sse_clients.append(q)

    # Send the current state immediately on connect so the page
    # renders QR codes even if no new update arrives yet.
    if current_urls.get('base_url'):
        initial = json.dumps({
            'base_url':  current_urls['base_url'],
            'timestamp': current_urls['last_updated'],
            'portals':   current_urls['portals'],
        })
        q.put_nowait(f"data: {initial}\n\n")

    def generate():
        try:
            while True:
                try:
                    msg = q.get(timeout=25)
                    yield msg
                except Exception:
                    # Send keep-alive comment every 25 s to prevent proxy timeouts
                    yield ": keepalive\n\n"
        except GeneratorExit:
            pass
        finally:
            with _sse_lock:
                if q in _sse_clients:
                    _sse_clients.remove(q)

    return Response(
        stream_with_context(generate()),
        mimetype='text/event-stream',
        headers={
            'Cache-Control': 'no-cache',
            'X-Accel-Buffering': 'no',   # nginx: disable buffering
            'Connection': 'keep-alive',
        }
    )


@app.route('/api/qr-data', methods=['GET'])
def qr_data():
    window_param = request.args.get('window')
    base_url     = current_urls.get('base_url', '')

    if not base_url:
        return jsonify({
            'error':   'base_url not configured yet',
            'hint':    'POST {"base_url": "https://your-app.com"} to /api/update-urls',
            'portals': [],
        }), 503

    portals = current_urls.get('portals') or build_portals(base_url)

    if window_param:
        try:
            wnum   = int(window_param)
            portal = next((p for p in portals if p['window'] == wnum), None)
            if not portal:
                return jsonify({'error': f'Window {wnum} not found'}), 404
            return jsonify({'success': True, 'portal': portal})
        except ValueError:
            return jsonify({'error': 'window must be an integer'}), 400

    return jsonify({
        'success':      True,
        'base_url':     base_url,
        'last_updated': current_urls.get('last_updated', ''),
        'portals':      portals,
        'total':        len(portals),
    })


@app.route('/api/windows', methods=['GET'])
def windows():
    return jsonify({
        'success': True,
        'windows': WINDOW_DEFINITIONS,
        'total':   len(WINDOW_DEFINITIONS),
    })


if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)