"""
Anthropic + OpenAI + Mistral API Proxy Server v4.0
-----------------------------------------------------
API keys live in environment — clients never see them.

Routes
  /anthropic/<path>         → https://api.anthropic.com/<path>
  /openai/<path>            → https://api.openai.com/<path>
  /mistral/<path>           → https://api.mistral.ai/<path>
  /api/openai/upload-batch  → multipart file-upload helper (OAI Batch)
  /api/mistral/upload-batch → multipart file-upload helper (Mistral Batch)
  /api/health               → JSON status for all providers
  /api/batch-results/<id>   → legacy alias (Anthropic batch NDJSON)
  /*                        → static file serving
"""

from flask import Flask, request, Response, send_from_directory, jsonify
import requests
import io
import os

app = Flask(__name__)

ANTHROPIC_BASE = "https://api.anthropic.com"
OPENAI_BASE    = "https://api.openai.com"
MISTRAL_BASE   = "https://api.mistral.ai"
ANT_VERSION    = "2023-06-01"

AGENT_BETA_PATHS = ("/v1/agents", "/v1/vaults", "/v1/memory", "/v1/organizations/tunnels")
STREAM_CTS = ("text/event-stream", "application/x-ndjson", "application/octet-stream")


def get_anthropic_key(): return os.environ.get("ANTHROPIC_API_KEY", "")
def get_openai_key():    return os.environ.get("OPENAI_API_KEY", "")
def get_mistral_key():   return os.environ.get("MISTRAL_API_KEY", "")


# ═══════════════════════════════════════════════════
# SHARED PROXY CORE
# ═══════════════════════════════════════════════════

def proxy_to(base_url, headers, path):
    target = f"{base_url}/{path}"
    if request.query_string:
        target += "?" + request.query_string.decode()

    body = request.get_data() if request.method in ("POST", "PUT", "PATCH") else None

    try:
        r = requests.request(
            method=request.method,
            url=target,
            headers=headers,
            data=body,
            stream=True,
            timeout=300,
        )
        ct = r.headers.get("Content-Type", "application/json")
        is_stream = any(s in ct for s in STREAM_CTS)

        def generate():
            for chunk in r.iter_content(chunk_size=8192):
                if chunk:
                    yield chunk

        resp_hdrs = {"Content-Type": ct, "Access-Control-Allow-Origin": "*"}
        if is_stream:
            resp_hdrs["Cache-Control"] = "no-cache"
        return Response(generate() if is_stream else r.content,
                        status=r.status_code, headers=resp_hdrs)
    except requests.Timeout:
        return jsonify({"error": "Request timed out"}), 504
    except Exception as e:
        return jsonify({"error": str(e)}), 502


@app.after_request
def add_cors_headers(response):
    response.headers["Access-Control-Allow-Origin"]  = "*"
    response.headers["Access-Control-Allow-Methods"] = "GET, POST, PUT, PATCH, DELETE, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type, Authorization, anthropic-beta, anthropic-version"
    return response


def cors_preflight():
    return Response("", status=204, headers={
        "Access-Control-Allow-Origin": "*",
        "Access-Control-Allow-Methods": "GET, POST, PUT, PATCH, DELETE, OPTIONS",
        "Access-Control-Allow-Headers": "Content-Type, Authorization, anthropic-beta, anthropic-version",
    })


# ═══════════════════════════════════════════════════
# ANTHROPIC PROXY
# ═══════════════════════════════════════════════════

def anthropic_headers(incoming, path=""):
    hdrs = {
        "x-api-key": get_anthropic_key(),
        "anthropic-version": ANT_VERSION,
    }
    for h in ("content-type", "anthropic-beta"):
        v = incoming.get(h)
        if v:
            hdrs[h] = v
    if "anthropic-beta" not in hdrs:
        for prefix in AGENT_BETA_PATHS:
            if path.startswith(prefix):
                hdrs["anthropic-beta"] = (
                    "mcp-tunnels-2026-05-19" if "tunnels" in prefix
                    else "managed-agents-2026-04-01"
                )
                break
    return hdrs


@app.route("/anthropic/<path:path>",
           methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"])
def anthropic_proxy(path):
    if request.method == "OPTIONS":
        return cors_preflight()
    if not get_anthropic_key():
        return jsonify({"error": "ANTHROPIC_API_KEY not configured"}), 500
    return proxy_to(ANTHROPIC_BASE, anthropic_headers(request.headers, f"/{path}"), path)


# ═══════════════════════════════════════════════════
# OPENAI PROXY
# ═══════════════════════════════════════════════════

def bearer_headers(key, incoming=None):
    hdrs = {"Authorization": f"Bearer {key}"}
    if incoming:
        ct = incoming.get("content-type")
        if ct:
            hdrs["content-type"] = ct
    return hdrs


@app.route("/openai/<path:path>",
           methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"])
def openai_proxy(path):
    if request.method == "OPTIONS":
        return cors_preflight()
    key = get_openai_key()
    if not key:
        return jsonify({"error": "OPENAI_API_KEY not configured"}), 500
    return proxy_to(OPENAI_BASE, bearer_headers(key, request.headers), path)


@app.route("/api/openai/upload-batch", methods=["POST"])
def openai_upload_batch():
    key = get_openai_key()
    if not key:
        return jsonify({"error": "OPENAI_API_KEY not configured"}), 500
    jsonl_bytes = request.get_data()
    if not jsonl_bytes:
        return jsonify({"error": "No JSONL data received"}), 400
    try:
        r = requests.post(
            f"{OPENAI_BASE}/v1/files",
            headers={"Authorization": f"Bearer {key}"},
            files={"file": ("batch_input.jsonl", io.BytesIO(jsonl_bytes), "application/x-ndjson")},
            data={"purpose": "batch"},
            timeout=120,
        )
        return Response(r.content, status=r.status_code,
                        headers={"Content-Type": "application/json",
                                 "Access-Control-Allow-Origin": "*"})
    except Exception as e:
        return jsonify({"error": str(e)}), 502


# ═══════════════════════════════════════════════════
# MISTRAL PROXY
# ═══════════════════════════════════════════════════

@app.route("/mistral/<path:path>",
           methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"])
def mistral_proxy(path):
    if request.method == "OPTIONS":
        return cors_preflight()
    key = get_mistral_key()
    if not key:
        return jsonify({"error": "MISTRAL_API_KEY not configured"}), 500
    return proxy_to(MISTRAL_BASE, bearer_headers(key, request.headers), path)


@app.route("/api/mistral/upload-batch", methods=["POST"])
def mistral_upload_batch():
    key = get_mistral_key()
    if not key:
        return jsonify({"error": "MISTRAL_API_KEY not configured"}), 500
    jsonl_bytes = request.get_data()
    if not jsonl_bytes:
        return jsonify({"error": "No JSONL data received"}), 400
    try:
        r = requests.post(
            f"{MISTRAL_BASE}/v1/files",
            headers={"Authorization": f"Bearer {key}"},
            files={"file": ("batch_input.jsonl", io.BytesIO(jsonl_bytes), "application/x-ndjson")},
            data={"purpose": "batch"},
            timeout=120,
        )
        return Response(r.content, status=r.status_code,
                        headers={"Content-Type": "application/json",
                                 "Access-Control-Allow-Origin": "*"})
    except Exception as e:
        return jsonify({"error": str(e)}), 502


# ═══════════════════════════════════════════════════
# HEALTH CHECK
# ═══════════════════════════════════════════════════

@app.route("/api/health")
def health():
    ant = bool(get_anthropic_key())
    oai = bool(get_openai_key())
    mis = bool(get_mistral_key())
    all_ok = ant and oai and mis
    return jsonify({
        "status": "ok" if all_ok else ("partial" if (ant or oai or mis) else "degraded"),
        "anthropic_key_configured": ant,
        "openai_key_configured": oai,
        "mistral_key_configured": mis,
        "proxy_targets": {
            "anthropic": ANTHROPIC_BASE,
            "openai": OPENAI_BASE,
            "mistral": MISTRAL_BASE,
        },
        "version": "4.0.0",
    })


# ═══════════════════════════════════════════════════
# BACKWARD COMPAT
# ═══════════════════════════════════════════════════

@app.route("/api/batch-results/<batch_id>")
def batch_results_compat(batch_id):
    if not get_anthropic_key():
        return jsonify({"error": "ANTHROPIC_API_KEY not configured"}), 500
    return proxy_to(ANTHROPIC_BASE,
                    anthropic_headers(request.headers),
                    f"v1/messages/batches/{batch_id}/results")


# ═══════════════════════════════════════════════════
# STATIC FILE SERVING
# ═══════════════════════════════════════════════════

@app.route("/", defaults={"path": "index.html"})
@app.route("/<path:path>")
def serve_static(path):
    if path.startswith(("api/", "anthropic/", "openai/", "mistral/")):
        return jsonify({"error": "Not found"}), 404
    try:
        return send_from_directory(".", path)
    except Exception:
        return send_from_directory(".", "index.html")


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    ant_s = "✓" if get_anthropic_key() else "✗ MISSING"
    oai_s = "✓" if get_openai_key()    else "✗ MISSING"
    mis_s = "✓" if get_mistral_key()   else "✗ MISSING"
    print(f"\n  ⚡ ANTHROPIC + OPENAI + MISTRAL PROXY v4.0")
    print(f"  ─────────────────────────────────────────")
    print(f"  Port          : {port}")
    print(f"  Anthropic Key : {ant_s}")
    print(f"  OpenAI Key    : {oai_s}")
    print(f"  Mistral Key   : {mis_s}")
    print(f"  /anthropic/<p> → api.anthropic.com")
    print(f"  /openai/<p>    → api.openai.com")
    print(f"  /mistral/<p>   → api.mistral.ai")
    print()
    app.run(host="0.0.0.0", port=port, threaded=True)