import os
import time
import json
import requests
import logging
from pathlib import Path

logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(message)s')
logger = logging.getLogger("wavespeed_client")

class WavespeedClient:
    def __init__(self, api_key=None):
        self.api_key = api_key or os.environ.get("WAVESPEED_API_KEY")
        if not self.api_key:
            raise ValueError("WAVESPEED_API_KEY not found in environment.")
        
        self.base_url = "https://api.wavespeed.ai/api/v3"
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

    def submit_task(self, model_endpoint, payload):
        """Submits a generation task to a specific model endpoint."""
        logger.info(f"Submitting task to {model_endpoint}...")
        url = f"{self.base_url}/{model_endpoint}"
        
        response = requests.post(url, headers=self.headers, json=payload)
        response.raise_for_status()
        
        result = response.json()
        task_id = result["data"]["id"]
        logger.info(f"Task submitted. ID: {task_id}")
        return task_id

    def poll_result(self, task_id, interval=10, timeout=600):
        """Polls for the result of a task until completion or timeout."""
        logger.info(f"Polling for result of task {task_id}...")
        url = f"{self.base_url}/predictions/{task_id}/result"
        
        start_time = time.time()
        while time.time() - start_time < timeout:
            response = requests.get(url, headers=self.headers)
            response.raise_for_status()
            
            result = response.json()
            data = result["data"]
            status = data.get("status")
            
            if status == "completed":
                outputs = data.get("outputs", [])
                logger.info(f"Task completed successfully. Outputs: {outputs}")
                return outputs
            elif status == "failed":
                error = data.get("error", "Unknown error")
                logger.error(f"Task failed: {error}")
                raise Exception(f"Wavespeed task failed: {error}")
            
            time.sleep(interval)
            
        raise TimeoutError(f"Task {task_id} timed out after {timeout} seconds.")