import subprocess
import time
import threading
import argilla as rg
import gradio as gr
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
import httpx

# Create FastAPI app
app = FastAPI(title="Argilla + Gradio Space")

# Add CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def run_argilla_server():
    """Run Argilla server in the background on port 6900"""
    subprocess.run([
        "uvicorn",
        "argilla.server:app",
        "--host", "0.0.0.0",
        "--port", "6900"
    ])

# Start Argilla in background thread
threading.Thread(target=run_argilla_server, daemon=True).start()

# ⏳ Wait longer for Argilla to boot
time.sleep(15)

# Initialize Argilla client (optional, just for verification)
try:
    rg.init(
        api_url="http://localhost:6900",
        api_key="admin.apikey",
        workspace="admin"
    )
except Exception as e:
    print("Argilla init failed:", e)

# ✅ Proxy /api/* to Argilla
@app.api_route("/api/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE"])
async def proxy_to_argilla(request: Request, path: str):
    url = f"http://localhost:6900/api/{path}"
    async with httpx.AsyncClient(follow_redirects=True) as client:
        resp = await client.request(
            method=request.method,
            url=url,
            headers={k: v for k, v in request.headers.items() if k.lower() != "host"},
            content=await request.body()
        )
    # Clean up headers that break FastAPI
    headers = {k: v for k, v in resp.headers.items() if k.lower() not in ["content-encoding", "transfer-encoding", "content-length"]}
    return StreamingResponse(resp.aiter_raw(), status_code=resp.status_code, headers=headers)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Argilla Server + Client")
    status_box = gr.Textbox("Server proxy at /api", label="Status")
    api_info = gr.Textbox("Try /api/health or /api/v1/me", label="API Info")
    
    def get_status():
        return "Server is running", "API endpoints: /api/health, /api/v1/me"
    
    btn = gr.Button("Check Status")
    out = gr.Textbox(label="Output")
    btn.click(fn=get_status, outputs=[status_box, out])

# Mount Gradio UI
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)