import asyncio
import os
import signal
import subprocess
import sys
from contextlib import asynccontextmanager

import httpx
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import Response, StreamingResponse

MANAGER_URL = "http://127.0.0.1:7894"
API_URL = "http://127.0.0.1:7895"

processes = []


def start_process(command):
    env = os.environ.copy()
    env.setdefault("DB_PATH", "/data/privy_manager.db")
    return subprocess.Popen(command, env=env)


@asynccontextmanager
async def lifespan(app: FastAPI):
    processes.append(start_process([sys.executable, "src/privy_manager.py"]))
    processes.append(start_process([sys.executable, "src/api_server.py"]))
    await asyncio.sleep(3)
    try:
        yield
    finally:
        for proc in processes:
            if proc.poll() is None:
                proc.send_signal(signal.SIGTERM)
        for proc in processes:
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                proc.kill()


app = FastAPI(lifespan=lifespan)


async def proxy(request: Request, upstream: str, path: str):
    url = f"{upstream}/{path}" if path else upstream
    if request.url.query:
        url = f"{url}?{request.url.query}"

    headers = dict(request.headers)
    headers.pop("host", None)
    body = await request.body()

    client = httpx.AsyncClient(timeout=None, follow_redirects=False)
    req = client.build_request(
        request.method,
        url,
        headers=headers,
        content=body,
    )
    upstream_resp = await client.send(req, stream=True)

    excluded = {"content-encoding", "transfer-encoding", "connection"}
    response_headers = {k: v for k, v in upstream_resp.headers.items() if k.lower() not in excluded}

    async def stream_response():
        try:
            async for chunk in upstream_resp.aiter_raw():
                yield chunk
        finally:
            await upstream_resp.aclose()
            await client.aclose()

    return StreamingResponse(
        stream_response(),
        status_code=upstream_resp.status_code,
        headers=response_headers,
        media_type=upstream_resp.headers.get("content-type"),
    )


@app.api_route("/v1/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def api_proxy(request: Request, path: str):
    return await proxy(request, API_URL, f"v1/{path}")


@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def manager_proxy(request: Request, path: str):
    return await proxy(request, MANAGER_URL, path)


if __name__ == "__main__":
    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)