scripts/cloud_endpoint.py

"""
COSMOS Cloud Endpoint API

REST API for cloud service integration and configuration.
Allows users to input API credentials and route requests to their cloud providers.
"""

from flask import Flask, request, jsonify, render_template
from cloud_router import CloudConfig, CloudRouter
import os
import json

app = Flask(__name__)
config = CloudConfig()
router = CloudRouter(config)

# ============================================================================
# CONFIGURATION ENDPOINTS
# ============================================================================

@app.route("/api/config", methods=["GET"])
def get_config():
    """Get current cloud configuration (sanitized, no keys)."""
    safe_config = json.loads(json.dumps(config.config))
    # Strip API keys
    for provider in safe_config["providers"].values():
        if "api_key" in provider:
            provider["api_key"] = "[HIDDEN]"
    
    return jsonify(safe_config)


@app.route("/api/config/provider/<provider>", methods=["GET"])
def get_provider_config(provider):
    """Get config for specific provider."""
    cfg = config.get_provider_config(provider)
    if not cfg:
        return jsonify({"error": f"Unknown provider: {provider}"}), 404
    
    # Sanitize
    if "api_key" in cfg:
        cfg["api_key"] = "[HIDDEN]"
    
    return jsonify(cfg)


@app.route("/api/config/provider/<provider>", methods=["POST"])
def update_provider_config(provider):
    """Update provider configuration."""
    data = request.json
    
    if provider not in config.config["providers"]:
        return jsonify({"error": f"Unknown provider: {provider}"}), 404
    
    # Update configuration
    if "endpoint" in data:
        config.set_provider_endpoint(provider, data["endpoint"])
    
    if "model_name" in data:
        model_blob = data.get("model_blob", data["model_name"])
        config.set_provider_model(provider, data["model_name"], model_blob)
    
    if "enabled" in data:
        config.enable_provider(provider, data["enabled"])
    
    # Save (without keys)
    config.save()
    
    return jsonify({"status": "updated", "provider": provider})


@app.route("/api/config/provider/<provider>/enable", methods=["POST"])
def enable_provider(provider):
    """Enable a provider."""
    config.enable_provider(provider, True)
    config.save()
    return jsonify({"status": "enabled", "provider": provider})


@app.route("/api/config/provider/<provider>/disable", methods=["POST"])
def disable_provider(provider):
    """Disable a provider."""
    config.enable_provider(provider, False)
    config.save()
    return jsonify({"status": "disabled", "provider": provider})


@app.route("/api/config/default", methods=["GET"])
def get_default_provider():
    """Get default provider."""
    return jsonify({"default_provider": config.get_active_provider()})


@app.route("/api/config/default", methods=["POST"])
def set_default_provider():
    """Set default provider."""
    data = request.json
    provider = data.get("provider")
    
    if provider not in config.config["providers"]:
        return jsonify({"error": f"Unknown provider: {provider}"}), 400
    
    config.set_default_provider(provider)
    config.save()
    
    return jsonify({"status": "default_set", "provider": provider})


# ============================================================================
# CREDENTIAL ENDPOINTS (ENVIRONMENT VARIABLE BASED)
# ============================================================================

@app.route("/api/credentials/azure", methods=["POST"])
def set_azure_credentials():
    """Set Azure credentials via environment variable."""
    data = request.json
    
    if "api_key" not in data:
        return jsonify({"error": "Missing api_key"}), 400
    
    # Set environment variable (in memory only, not persisted)
    os.environ["COSMOS_AZURE_KEY"] = data["api_key"]
    config.load_credentials_from_env()
    
    return jsonify({"status": "credentials_set", "provider": "azure"})


@app.route("/api/credentials/ibm", methods=["POST"])
def set_ibm_credentials():
    """Set IBM credentials via environment variable."""
    data = request.json
    
    if "api_key" not in data:
        return jsonify({"error": "Missing api_key"}), 400
    
    # Set environment variable (in memory only, not persisted)
    os.environ["COSMOS_IBM_KEY"] = data["api_key"]
    config.load_credentials_from_env()
    
    return jsonify({"status": "credentials_set", "provider": "ibm"})


@app.route("/api/credentials/test/<provider>", methods=["POST"])
def test_credentials(provider):
    """Test provider credentials and connectivity."""
    try:
        response = router.generate("Hello", provider=provider)
        return jsonify({"status": "success", "provider": provider, "sample": response[:100]})
    except Exception as e:
        return jsonify({"status": "error", "provider": provider, "error": str(e)}), 400


# ============================================================================
# GENERATION ENDPOINTS
# ============================================================================

@app.route("/api/generate", methods=["POST"])
def generate():
    """Generate response from configured provider."""
    data = request.json
    prompt = data.get("prompt")
    provider = data.get("provider")  # Optional, uses default if not provided
    
    if not prompt:
        return jsonify({"error": "Missing prompt"}), 400
    
    try:
        response = router.generate(prompt, provider=provider)
        return jsonify({"response": response, "provider": provider or config.get_active_provider()})
    except Exception as e:
        return jsonify({"error": str(e)}), 400


@app.route("/api/vision", methods=["POST"])
def vision():
    """Analyze image with vision model."""
    # Expects multipart form with 'image' file and 'prompt' field
    
    if "image" not in request.files:
        return jsonify({"error": "Missing image file"}), 400
    
    prompt = request.form.get("prompt", "Describe this image.")
    provider = request.form.get("provider")
    
    image_file = request.files["image"]
    image_path = f"/tmp/{image_file.filename}"
    image_file.save(image_path)
    
    try:
        response = router.vision(image_path, prompt, provider=provider)
        return jsonify({"response": response, "provider": provider or config.get_active_provider()})
    except Exception as e:
        return jsonify({"error": str(e)}), 400
    finally:
        if os.path.exists(image_path):
            os.remove(image_path)


# ============================================================================
# STATUS & HEALTH ENDPOINTS
# ============================================================================

@app.route("/api/status", methods=["GET"])
def status():
    """Get system status."""
    providers = config.config["providers"]
    status_info = {
        "default_provider": config.get_active_provider(),
        "providers": {
            name: {
                "enabled": cfg.get("enabled", False),
                "endpoint": cfg.get("api_endpoint", "N/A"),
                "model": cfg.get("model_name", "N/A"),
                "has_credentials": bool(cfg.get("api_key"))
            }
            for name, cfg in providers.items()
        }
    }
    return jsonify(status_info)


@app.route("/api/health", methods=["GET"])
def health():
    """Health check."""
    return jsonify({"status": "healthy", "service": "COSMOS Cloud Router"})


# ============================================================================
# UI ENDPOINTS
# ============================================================================

@app.route("/", methods=["GET"])
def dashboard():
    """Configuration dashboard (simple HTML form)."""
    return render_template("dashboard.html")


@app.route("/docs", methods=["GET"])
def docs():
    """API documentation."""
    return render_template("api_docs.html")


# ============================================================================
# ERROR HANDLERS
# ============================================================================

@app.errorhandler(404)
def not_found(error):
    return jsonify({"error": "Endpoint not found"}), 404


@app.errorhandler(500)
def internal_error(error):
    return jsonify({"error": "Internal server error"}), 500


# ============================================================================
# MAIN
# ============================================================================

if __name__ == "__main__":
    print("COSMOS Cloud Endpoint Starting...")
    print(f"Default provider: {config.get_active_provider()}")
    print(f"Enabled providers: {[name for name, cfg in config.config['providers'].items() if cfg.get('enabled')]}")
    print("\nAPI Documentation: http://localhost:5000/docs")
    print("Configuration Dashboard: http://localhost:5000/")
    print("\nEnvironment variables for credentials:")
    print("  COSMOS_AZURE_KEY  - Azure OpenAI API key")
    print("  COSMOS_IBM_KEY    - IBM Watsonx API key")
    
    app.run(debug=True, host="0.0.0.0", port=5000)