import os
from typing import Optional, Dict, Any
import httpx
from pydantic import BaseModel

class OptimizationStats(BaseModel):
    original_tokens: int
    optimized_tokens: int
    reduction_percent: int
    latency_ms: float

class OptimizationResult(BaseModel):
    text: str
    stats: OptimizationStats

class Trunkate:
    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        self.api_key = api_key or os.getenv("TRUNKATE_API_KEY")
        if not self.api_key:
            raise ValueError(
                "The TRUNKATE_API_KEY environment variable is missing or empty; "
                "either provide it, or instantiate the Trunkate client with an api_key."
            )
        # Default to official API if not specified, can be overridden by env for testing
        self.base_url = base_url or os.getenv("TRUNKATE_API_URL", "https://api.trunkate.ai")
        # Ensure trailing slashes are stripped
        self.base_url = self.base_url.rstrip("/")
        
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        self.client = httpx.Client(base_url=self.base_url, headers=headers, timeout=10.0)

    def optimize(self, text: str, task: Optional[str] = None, budget: Optional[int] = None, model: str = "gpt-4o") -> OptimizationResult:
        """
        Optimize a prompt using the Trunkate AI API.
        """
        payload = {
            "text": text,
            "model": model
        }
        if task:
            payload["task"] = task
        if budget:
            payload["budget"] = budget

        response = self.client.post("/optimize", json=payload)
        response.raise_for_status()
        
        data = response.json()
        
        return OptimizationResult(
            text=data["optimized_text"],
            stats=OptimizationStats(
                original_tokens=data["original_tokens"],
                optimized_tokens=data["optimized_tokens"],
                reduction_percent=data["reduction_percent"],
                latency_ms=data["latency_ms"]
            )
        )
