frontend/api_client.py

import os
import requests
from typing import List, Dict, Any, Optional, Tuple

from dotenv import load_dotenv

# Load environment variables from backend/.env if it exists
backend_env = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "backend", ".env")
load_dotenv(backend_env)

BACKEND_URL = os.environ.get("BACKEND_URL")
if not BACKEND_URL:
    if "SPACE_ID" in os.environ:
        # In Hugging Face Spaces, the backend runs internally on 127.0.0.1:8000
        BACKEND_URL = "http://127.0.0.1:8000"
    else:
        backend_host = os.environ.get("HOST", "127.0.0.1")
        backend_port = os.environ.get("PORT", "8000")
        BACKEND_URL = f"http://{backend_host}:{backend_port}"

class AntigravityAPIClient:
    """
    HTTP Client that routes queries from Streamlit UI to the stateless FastAPI backend.
    """
    def __init__(self, base_url: str = BACKEND_URL):
        self.base_url = base_url.rstrip("/")

    def upload_files(self, file_tuples: List[Tuple[str, bytes]]) -> Dict[str, Any]:
        """
        Uploads up to 10 Excel files to the backend.
        file_tuples: list of (filename, file_bytes)
        """
        url = f"{self.base_url}/api/upload"
        files = []
        for name, data in file_tuples:
            if name.lower().endswith(".csv"):
                mime = "text/csv"
            else:
                mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
            files.append(("files", (name, data, mime)))
            
        try:
            response = requests.post(url, files=files, timeout=120)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Upload failed: {str(e)}"}

    def filter_dataset(
        self,
        filepath: str,
        sheet_name: str,
        filters: List[Dict[str, Any]],
        tree_group_cols: Optional[List[str]] = None,
        cluster_cols: Optional[List[str]] = None,
        cluster_count: int = 3
    ) -> Dict[str, Any]:
        """
        Requests universal filtering, clustering, and tri-view representations.
        """
        url = f"{self.base_url}/api/filter"
        payload = {
            "filepath": filepath,
            "sheet_name": sheet_name,
            "filters": filters,
            "tree_group_cols": tree_group_cols or [],
            "cluster_cols": cluster_cols or [],
            "cluster_count": cluster_count
        }
        try:
            response = requests.post(url, json=payload, timeout=60)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            err_msg = response.json().get("detail", str(e)) if 'response' in locals() and response else str(e)
            return {"error": f"Filtering failed: {err_msg}"}

    def execute_pivot(
        self,
        filepath: str,
        sheet_name: str,
        index: List[str],
        columns: List[str],
        values: str,
        agg: str
    ) -> Dict[str, Any]:
        """
        Requests custom pivot aggregation.
        """
        url = f"{self.base_url}/api/pivot"
        payload = {
            "filepath": filepath,
            "sheet_name": sheet_name,
            "index": index,
            "columns": columns,
            "values": values,
            "agg": agg
        }
        try:
            response = requests.post(url, json=payload, timeout=60)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            err_msg = response.json().get("detail", str(e)) if 'response' in locals() and response else str(e)
            return {"error": f"Pivot failed: {err_msg}"}

    def execute_join(
        self,
        filepath_a: str,
        sheet_name_a: str,
        filepath_b: str,
        sheet_name_b: str,
        join_key_a: str,
        join_key_b: str,
        join_type: str,
        select_columns_a: List[str],
        select_columns_b: List[str]
    ) -> Any:
        """
        Triggers vectorized relational join and returns binary Excel stream.
        """
        url = f"{self.base_url}/api/join"
        payload = {
            "filepath_a": filepath_a,
            "sheet_name_a": sheet_name_a,
            "filepath_b": filepath_b,
            "sheet_name_b": sheet_name_b,
            "join_key_a": join_key_a,
            "join_key_b": join_key_b,
            "join_type": join_type.lower(),
            "select_columns_a": select_columns_a,
            "select_columns_b": select_columns_b
        }
        try:
            response = requests.post(url, json=payload, stream=True, timeout=120)
            response.raise_for_status()
            joined_filepath = response.headers.get("X-Joined-Filepath", "")
            return {
                "content": response.content,
                "joined_filepath": joined_filepath
            }
        except requests.exceptions.RequestException as e:
            err_msg = response.text if 'response' in locals() and response else str(e)
            return {"error": f"Join failed: {err_msg}"}

    def get_diagnostics(self) -> Dict[str, Any]:
        """
        Fetches backend hardware and optimization state.
        """
        url = f"{self.base_url}/api/diagnostics"
        try:
            response = requests.get(url, timeout=5)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {"error": f"Diagnostics retrieval failed: {str(e)}"}

    def ask_analyst(self, filepath: str, sheet_name: str, question: str) -> Dict[str, Any]:
        """
        Asks LangChain AI Data Analyst a semantic question.
        """
        url = f"{self.base_url}/api/chat"
        payload = {
            "filepath": filepath,
            "sheet_name": sheet_name,
            "question": question
        }
        try:
            response = requests.post(url, json=payload, timeout=60)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            err_msg = response.json().get("detail", str(e)) if 'response' in locals() and response else str(e)
            return {"error": f"LangChain semantic lookup failed: {err_msg}"}

    def update_remarks(self, filepath: str, sheet_name: str, updates: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Saves user-modified Remarks back into the Excel/CSV file on the server.
        """
        url = f"{self.base_url}/api/update_remarks"
        payload = {
            "filepath": filepath,
            "sheet_name": sheet_name,
            "updates": updates
        }
        try:
            response = requests.post(url, json=payload, timeout=60)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            err_msg = response.json().get("detail", str(e)) if 'response' in locals() and response else str(e)
            return {"error": f"Updating remarks failed: {err_msg}"}