from typing import Optional, Dict, Any
import os
import httpx
from .types import AddResponse, SearchResponse


class Client:
    
    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: str = "https://api.lithio.dev",
        timeout: float = 30.0,
    ):
        self.api_key = api_key or os.getenv("LITHIO_API_KEY")
        if not self.api_key:
            raise ValueError("api_key must be provided or LITHIO_API_KEY environment variable must be set")
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._client = httpx.Client(
            base_url=self.base_url,
            timeout=timeout,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
        )
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
    
    def close(self):
        self._client.close()
    
    def add(
        self,
        query: str,
        space: Optional[str] = None,
        properties: Optional[Dict[str, Any]] = None,
    ) -> AddResponse:
        payload = {
            "query": query,
        }
        
        if space is not None:
            payload["space"] = space
        
        if properties:
            payload["properties"] = properties
        
        response = self._client.post("/add", json=payload)
        response.raise_for_status()
        return AddResponse(**response.json())
    
    def search(
        self,
        query: str,
        limit: Optional[int] = None,
        space: Optional[str] = None,
        properties: Optional[Dict[str, Any]] = None,
    ) -> SearchResponse:
        payload = {
            "query": query,
        }
        
        if limit is not None:
            payload["limit"] = limit
        
        if space is not None:
            payload["space"] = space
        
        if properties:
            payload["properties"] = properties
        
        response = self._client.post("/search", json=payload)
        response.raise_for_status()
        return SearchResponse(**response.json())

