client.py

"""Dynamic OpenAPI Client Driver for SOG-X Backend."""
import mimetypes
import os
import tempfile
import logging
from typing import Dict, Any, Tuple, Optional
import httpx

logger = logging.getLogger("sogx.client")


class BackendClient:
    def __init__(self, backend_url: Optional[str] = None, token: Optional[str] = None):
        self.base_url = (backend_url or os.getenv("BACKEND_URL", "http://localhost:7860")).rstrip("/")
        self.token = token or os.getenv("HF_TOKEN") or os.getenv("API_KEY", "")

    def get_headers(self, user_api_key: Optional[str] = None) -> dict:
        headers = {}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        active_key = user_api_key or os.getenv("USER_KEY", "")
        if active_key:
            headers["X-API-Key"] = active_key
            headers["X-User-API-Key"] = active_key
        return headers

    def fetch_openapi_schema(self, user_api_key: Optional[str] = None) -> Tuple[bool, Dict[str, Any]]:
        """Auto-discover backend endpoints by pulling live openapi.json."""
        url = f"{self.base_url}/openapi.json"
        try:
            with httpx.Client(timeout=10.0) as client:
                res = client.get(url, headers=self.get_headers(user_api_key))
                if res.status_code == 200:
                    return True, res.json()
                return False, {"error": f"HTTP {res.status_code}", "detail": res.text}
        except Exception as exc:
            return False, {"error": "NETWORK_ERROR", "detail": str(exc)}

    def execute_dynamic_request(
        self,
        path: str,
        method: str,
        params: Dict[str, Any],
        payload: Dict[str, Any],
        files: Dict[str, Any],
        user_api_key: Optional[str] = None,
    ) -> Tuple[bool, Any]:
        """Execute request against any backend route dynamically.

        On success, returns (True, response) where response is one of:
        - a parsed JSON dict/list (for application/json responses)
        - {"__binary__": True, "content_type": str, "file_path": str} for
          images, audio, video, PDFs, or any other downloadable file response
        - a plain string for anything else (e.g. text/plain)
        """
        url = f"{self.base_url}{path}"
        headers = self.get_headers(user_api_key)

        try:
            with httpx.Client(timeout=300.0) as client:
                if method.upper() == "GET":
                    res = client.get(url, headers=headers, params=params)
                elif method.upper() == "POST":
                    if files:
                        res = client.post(url, headers=headers, params=params, data=payload, files=files)
                    else:
                        res = client.post(url, headers=headers, params=params, json=payload)
                else:
                    return False, f"Unsupported HTTP method: {method}"

                if res.status_code == 200:
                    return True, self._parse_response(res)
                elif res.status_code == 401:
                    return False, "❌ 401 Unauthorized: Invalid API Key."
                elif res.status_code == 429:
                    return False, "⚠️ 429 Rate Limited: Daily quota reached."
                return False, f"Error {res.status_code}: {res.text}"

        except Exception as exc:
            return False, f"Request Failed: {str(exc)}"

    @staticmethod
    def _parse_response(res: "httpx.Response") -> Any:
        """Route the response body based on its real Content-Type, the way Swagger UI does."""
        content_type = res.headers.get("content-type", "").split(";")[0].strip().lower()
        disposition = res.headers.get("content-disposition", "")

        if content_type.startswith("application/json"):
            try:
                return res.json()
            except Exception:
                pass  # fall through and treat as text/binary below

        is_binary = (
            content_type.startswith(("image/", "audio/", "video/", "application/pdf"))
            or content_type == "application/octet-stream"
            or "attachment" in disposition
        )
        if is_binary:
            ext = mimetypes.guess_extension(content_type) or ""
            fd, tmp_path = tempfile.mkstemp(suffix=ext)
            with os.fdopen(fd, "wb") as fh:
                fh.write(res.content)
            return {"__binary__": True, "content_type": content_type, "file_path": tmp_path}

        try:
            return res.json()
        except Exception:
            return res.text