from __future__ import annotations

import os
from collections.abc import AsyncIterator
from typing import Annotated

import httpx
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
from open_webui.utils.auth import get_verified_user
from open_webui.utils.headers import include_user_info_headers

router = APIRouter()

GATEWAY_URL = os.getenv("K1412_GATEWAY_URL", "http://gateway:8001").rstrip("/")
GATEWAY_KEY = os.getenv("K1412_GATEWAY_KEY", "")
REQUEST_TIMEOUT = httpx.Timeout(120, connect=10)
STREAM_TIMEOUT = httpx.Timeout(1800, connect=10)
VerifiedUser = Annotated[object, Depends(get_verified_user)]


def gateway_headers(user) -> dict[str, str]:
    if not GATEWAY_KEY:
        raise HTTPException(status_code=503, detail="Workspace service is not configured")
    return include_user_info_headers(
        {"Authorization": f"Bearer {GATEWAY_KEY}"},
        user,
    )


async def gateway_error(response: httpx.Response) -> HTTPException:
    try:
        payload = response.json()
        detail = payload.get("detail", payload)
    except Exception:
        detail = (await response.aread()).decode("utf-8", errors="replace")[:2000]
    return HTTPException(status_code=response.status_code, detail=detail)


@router.get("/files")
async def list_workspace_files(
    user: VerifiedUser,
    path: str = Query(default=".", max_length=4096),
):
    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        response = await client.get(
            f"{GATEWAY_URL}/v1/files",
            params={"path": path},
            headers=gateway_headers(user),
        )
        if response.status_code >= 400:
            raise await gateway_error(response)
        return response.json()


async def close_after_stream(response: httpx.Response, client: httpx.AsyncClient) -> AsyncIterator[bytes]:
    try:
        async for chunk in response.aiter_bytes():
            yield chunk
    finally:
        await response.aclose()
        await client.aclose()


async def stream_workspace_response(endpoint: str, path: str, user) -> StreamingResponse:
    client = httpx.AsyncClient(timeout=STREAM_TIMEOUT)
    request = client.build_request(
        "GET",
        f"{GATEWAY_URL}{endpoint}",
        params={"path": path},
        headers=gateway_headers(user),
    )
    response = await client.send(request, stream=True)
    if response.status_code >= 400:
        error = await gateway_error(response)
        await response.aclose()
        await client.aclose()
        raise error
    headers = {
        key: value
        for key, value in response.headers.items()
        if key.lower() in {"content-disposition", "content-length"}
    }
    return StreamingResponse(
        close_after_stream(response, client),
        media_type=response.headers.get("content-type", "application/octet-stream"),
        headers=headers,
    )


@router.get("/download")
async def download_workspace_file(
    user: VerifiedUser,
    path: str = Query(min_length=1, max_length=4096),
):
    return await stream_workspace_response("/v1/files/download", path, user)


@router.get("/archive")
async def archive_workspace_files(
    user: VerifiedUser,
    path: str = Query(default=".", max_length=4096),
):
    return await stream_workspace_response("/v1/files/archive", path, user)