import os

class Cie:
    def __init__(self, client):
        self.client = client

    def upload_dataset(self, file_path):
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found: {file_path}")
            
        with open(file_path, 'rb') as f:
            files = {'file': f}
            # Note: Requests handles multipart headers automatically when 'files' is passed
            # DO NOT set Content-Type: application/json manually here
            # We bypass the client wrapper to handle files properly
            url = f"{self.client.base_url}/cie/datasets"
            headers = self.client.session.headers.copy()
            if "Content-Type" in headers:
                del headers["Content-Type"] # Let requests set the boundary
            
            response = requests.post(url, files=files, headers=headers)
            return response.json()

    def list_datasets(self):
        return self.client.request("GET", "/cie/datasets")

    def train_model(self, dataset_id, name, model_type):
        return self.client.request("POST", "/cie/models/train", json={
            "datasetId": dataset_id,
            "name": name,
            "type": model_type
        })

    def list_models(self):
        return self.client.request("GET", "/cie/models")

    def predict(self, model_id, input_data):
        return self.client.request("POST", "/cie/predictions", json={
            "modelId": model_id,
            "input": input_data
        })
