import os
import time
import requests
from fastapi import FastAPI, UploadFile, File, Form, Query, Request, Body
from fastapi.responses import Response, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.concurrency import run_in_threadpool # <--- FIX: Added Threadpool
import uvicorn

app = FastAPI()

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

HF_TOKEN = os.getenv("HF_TOKEN")

DEV_URL = "https://PPAL-SongLab-UGA-fruit-analyzer-dev.hf.space"
PROD_URL = "https://PPAL-SongLab-UGA-fruit-analyzer.hf.space"
HF_RATE_LIMIT_STATUS = 429
HF_TRANSIENT_STATUS = {429, 502, 503, 504}
PROCESS_PROXY_TIMEOUT_SECONDS = 40


def request_with_hf_backoff(method, url, *, max_retries=3, **kwargs):
    last_response = None
    for attempt in range(max_retries + 1):
        response = method(url, **kwargs)
        last_response = response
        if response.status_code not in HF_TRANSIENT_STATUS:
            return response

        retry_after = response.headers.get("retry-after")
        try:
            wait_s = float(retry_after) if retry_after else None
        except ValueError:
            wait_s = None
        if wait_s is None:
            wait_s = [2.0, 5.0, 10.0, 15.0][min(attempt, 3)] if response.status_code == HF_RATE_LIMIT_STATUS else [3.0, 8.0, 15.0, 20.0][min(attempt, 3)]
        if attempt < max_retries:
            time.sleep(min(max(wait_s, 0.5), 20.0))
    return last_response


def hf_error_message(exc):
    response = getattr(exc, "response", None)
    if response is not None and response.status_code == HF_RATE_LIMIT_STATUS:
        retry_after = response.headers.get("retry-after")
        suffix = f" Retry after about {retry_after} seconds." if retry_after else " Please wait a minute and try again."
        return f"Hugging Face is rate limiting this Space after several retry attempts.{suffix}"
    if response is not None and response.status_code in {502, 503, 504}:
        return "Hugging Face Space is temporarily unavailable or waking up after several retry attempts. Please wait a moment and try again."
    if isinstance(exc, requests.exceptions.Timeout):
        return "Hugging Face Space did not respond before the proxy timeout. Please try again."
    if isinstance(exc, requests.exceptions.ConnectionError):
        return "Could not connect to the Hugging Face Space. Please try again."
    return str(exc)


def is_retryable_hf_exception(exc):
    response = getattr(exc, "response", None)
    return bool(response is not None and response.status_code in HF_TRANSIENT_STATUS)


def proxy_error_payload(prefix, exc, **extra):
    return {
        "success": False,
        "message": f"{prefix}: {hf_error_message(exc)}",
        "retryable": is_retryable_hf_exception(exc),
        **extra,
    }

@app.get("/")
def read_root():
    return {"status": "Render Proxy is awake!"}

@app.get("/proxy_status")
async def proxy_status(username: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    status_url = f"{base_url}/queue_status"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    
    # FIX: Define the blocking request
    def fetch_status():
        return requests.get(status_url, headers=headers, timeout=5)
        
    try:
        # FIX: Run it in a background thread so the server doesn't freeze
        response = await run_in_threadpool(fetch_status)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"active_requests": 0, "max_concurrent": 2}


@app.post("/proxy_warmup")
async def proxy_warmup(payload: dict = Body(...)):
    """Warm the production backend only.

    This deliberately does not route devtest to the dev Space; it exists to
    keep the production user path responsive without burning analysis jobs.
    """
    target_url = f"{PROD_URL}/warmup"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            json=payload,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Warmup Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Warmup Error (Internal): {str(e)}"}


@app.post("/proxy_process")
async def proxy_process(request: Request, file: UploadFile = File(...), password: str = Form(""), username: str = Form("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/process_single"

    query_str = request.url.query
    if query_str: target_url += f"?{query_str}"

    file_bytes = await file.read()
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    files = {"file": (file.filename, file_bytes, file.content_type)}
    form = await request.form()
    data = {}
    for key, value in form.multi_items():
        if key == "file" or hasattr(value, "filename"):
            continue
        data[key] = str(value)
    data.setdefault("password", password)
    data.setdefault("username", username)

    # FIX: Define the blocking request
    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            files=files,
            data=data,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        # FIX: Run it in a background thread so other users can check the status!
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.post("/proxy_compatibility")
async def proxy_compatibility(request: Request, password: str = Form(""), username: str = Form("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/compatibility_check"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    form = await request.form()
    files = []
    data = {}
    for key, value in form.multi_items():
        if hasattr(value, "filename"):
            files.append(("files", (value.filename, await value.read(), value.content_type)))
        else:
            data[key] = str(value)
    data.setdefault("password", password)
    data.setdefault("username", username)

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            files=files,
            data=data,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, results=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "results": []}


@app.get("/proxy_experts")
async def proxy_experts(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/experts"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, experts=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "experts": []}


@app.get("/proxy_training_model_options")
async def proxy_training_model_options(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/training_model_options"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, options=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "options": []}


@app.post("/proxy_finetune")
async def proxy_finetune(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    expert_id = str(payload.get("expert_id", "")).strip()
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/experts/{expert_id}/finetune"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(requests.post, target_url, max_retries=0, headers=headers, json=payload, timeout=60)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_train_job")
async def proxy_train_job(job_id: str = Query(...), username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/train_jobs/status/{job_id}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_train_jobs")
async def proxy_train_jobs(username: str = Query(""), password: str = Query("")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/train_jobs"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(target_url, headers=headers, params={"username": username, "password": password}, timeout=20)

    try:
        response = await run_in_threadpool(make_get)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, jobs=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "jobs": []}


DATASET_PROXY_TIMEOUT_SECONDS = 180


@app.api_route("/proxy_datasets{rest:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy_datasets(rest: str, request: Request):
    """Generic forwarder for all Labeling Studio dataset endpoints.

    Mirrors /proxy_datasets{rest} -> {base}/datasets{rest}, preserving the query
    string, body (multipart / JSON), and returning either JSON or raw binary
    (image/zip) responses. Routes dev vs prod by username (query, form, or JSON).
    """
    username = request.query_params.get("username", "")
    content_type = request.headers.get("content-type", "")

    files = []
    data = {}
    json_body = None
    if request.method in ("POST", "PUT", "DELETE"):
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
            username = username or data.get("username", "")
        elif "application/json" in content_type:
            try:
                json_body = await request.json()
            except Exception:
                json_body = None
            if isinstance(json_body, dict):
                username = username or str(json_body.get("username", ""))

    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/datasets{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body if json_body is not None else None,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
    except requests.exceptions.RequestException as e:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", e))
    except Exception as e:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Error (Internal): {str(e)}"})

    resp_ct = response.headers.get("content-type", "application/octet-stream")
    passthrough_headers = {}
    disposition = response.headers.get("content-disposition")
    if disposition:
        passthrough_headers["content-disposition"] = disposition
    return Response(
        content=response.content,
        status_code=response.status_code,
        media_type=resp_ct,
        headers=passthrough_headers,
    )

@app.post("/proxy_batch_stage")
async def proxy_batch_stage(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/batch_stage"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return request_with_hf_backoff(
            requests.post,
            target_url,
            max_retries=0,
            headers=headers,
            json=payload,
            timeout=PROCESS_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e, rows=[])
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}", "rows": []}


@app.api_route("/proxy_process_jobs{rest:path}", methods=["GET", "POST"])
async def proxy_process_jobs(rest: str, request: Request):
    """Forward persistent processing-job uploads, polling, stop, and resume."""
    try:
        username = request.query_params.get("username", "")
        content_type = request.headers.get("content-type", "")
        files = []
        data = {}
        json_body = None
        if request.method == "POST":
            if "multipart/form-data" in content_type:
                form = await request.form()
                for key, value in form.multi_items():
                    if hasattr(value, "filename"):
                        files.append((key, (value.filename, await value.read(), value.content_type)))
                    else:
                        data[key] = str(value)
                username = username or data.get("username", "")
            elif "application/json" in content_type:
                json_body = await request.json()
                if isinstance(json_body, dict):
                    username = username or str(json_body.get("username", ""))
    except Exception as exc:
        return JSONResponse(status_code=400, content={"success": False, "message": f"Proxy Process Job Error (Request): {str(exc)}"})
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/process_jobs{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=2,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
        return Response(
            content=response.content,
            status_code=response.status_code,
            media_type=response.headers.get("content-type", "application/json"),
        )
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))
    except Exception as exc:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Process Job Error (Internal): {str(exc)}"})


@app.post("/proxy_flush_queue")
async def proxy_flush_queue(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    if username.strip().lower() != "devtest":
        return {"success": False, "message": "Queue flush is only available for devtest."}
    target_url = f"{DEV_URL}/flush_queue"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return requests.post(target_url, headers=headers, json=payload, timeout=20)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        return proxy_error_payload("Proxy Error (Hugging Face)", e)
    except Exception as e:
        return {"success": False, "message": f"Proxy Error (Internal): {str(e)}"}


@app.get("/proxy_preview/{session_id}/{row_id}/{preview_type}")
async def proxy_preview(session_id: str, row_id: str, preview_type: str, username: str = Query(""), password: str = Query(""), size: str = Query("thumb")):
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/preview/{session_id}/{row_id}/{preview_type}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_get():
        return requests.get(
            target_url,
            headers=headers,
            params={"size": size, "username": username, "password": password},
            timeout=30,
        )

    response = await run_in_threadpool(make_get)
    if response.status_code >= 400:
        return Response(content=response.content, status_code=response.status_code, media_type=response.headers.get("content-type", "text/plain"))
    return Response(content=response.content, media_type=response.headers.get("content-type", "image/jpeg"))


@app.post("/proxy_preview_session_clear")
async def proxy_preview_session_clear(payload: dict = Body(...)):
    username = str(payload.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/preview_session/clear"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_post():
        return requests.post(target_url, headers=headers, json=payload, timeout=20)

    try:
        response = await run_in_threadpool(make_post)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"success": False, "message": f"Proxy Error: {str(e)}"}


@app.api_route("/proxy_adjustments{rest:path}", methods=["GET", "POST", "PUT"])
async def proxy_adjustments(rest: str, request: Request):
    """Forward preview-mask and traditional-feature adjustment requests."""
    username = request.query_params.get("username", "")
    content_type = request.headers.get("content-type", "")
    files = []
    data = {}
    json_body = None
    if request.method in {"POST", "PUT"}:
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
            username = username or data.get("username", "")
        elif "application/json" in content_type:
            json_body = await request.json()
            if isinstance(json_body, dict):
                username = username or str(json_body.get("username", ""))
    base_url = DEV_URL if username.strip().lower() == "devtest" else PROD_URL
    target_url = f"{base_url}/adjustments{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=DATASET_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
        return Response(
            content=response.content,
            status_code=response.status_code,
            media_type=response.headers.get("content-type", "application/json"),
        )
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))


SYNC_PROXY_TIMEOUT_SECONDS = 300


@app.api_route("/proxy_sync{rest:path}", methods=["GET", "POST"])
async def proxy_sync(rest: str, request: Request):
    """Forward admin sync requests to the dev backend only.

    Dev remains the controller for dev<->prod sync. The browser never receives
    production sync tokens, and this proxy never routes sync calls to prod based
    on username.
    """
    target_url = f"{DEV_URL}/sync{rest}"
    if request.url.query:
        target_url += f"?{request.url.query}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    content_type = request.headers.get("content-type", "")
    files = []
    data = {}
    json_body = None
    if request.method == "POST":
        if "multipart/form-data" in content_type:
            form = await request.form()
            for key, value in form.multi_items():
                if hasattr(value, "filename"):
                    files.append((key, (value.filename, await value.read(), value.content_type)))
                else:
                    data[key] = str(value)
        elif "application/json" in content_type:
            try:
                json_body = await request.json()
            except Exception:
                json_body = None

    def make_request():
        return request_with_hf_backoff(
            getattr(requests, request.method.lower()),
            target_url,
            max_retries=0,
            headers=headers,
            files=files or None,
            data=data or None,
            json=json_body,
            timeout=SYNC_PROXY_TIMEOUT_SECONDS,
        )

    try:
        response = await run_in_threadpool(make_request)
    except requests.exceptions.RequestException as exc:
        return JSONResponse(status_code=502, content=proxy_error_payload("Proxy Error (Hugging Face)", exc))
    except Exception as exc:
        return JSONResponse(status_code=500, content={"success": False, "message": f"Proxy Error (Internal): {str(exc)}"})

    media_type = response.headers.get("content-type", "application/json")
    passthrough_headers = {}
    disposition = response.headers.get("content-disposition")
    if disposition:
        passthrough_headers["content-disposition"] = disposition
    return Response(content=response.content, status_code=response.status_code, media_type=media_type, headers=passthrough_headers)


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