import httpx


class Client:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://api.onepot.ai"
        self.client = httpx.Client(
            headers={"X-API-Key": api_key},
            timeout=240.0,
            follow_redirects=True,
        )

    def search(
        self,
        smiles_list: list[str],
        max_results: int = 100,
        substructure_search: bool = False,
        max_depth: int = 1,
        include_chemistry_risk: bool = False,
        include_chemistry_risk_score: bool = False,
    ):
        """Search for purchasable analogs of the given molecules.

        Args:
            smiles_list: List of SMILES strings to search.
            max_results: Maximum results per query (default 100).
            substructure_search: Match substructures instead of similarity.
            max_depth: Synthesis depth (currently only 1 supported).
            include_chemistry_risk: Include chemistry_risk field (low/medium/high).
            include_chemistry_risk_score: Include raw probability score (0-1).

        Returns:
            Dict with:
                - queries: List of query results, each containing query_smiles, query_inchikey,
                  and results (list of {smiles, inchikey, similarity, supplier_risk, price_usd,
                  and optionally chemistry_risk/chemistry_risk_score}).
                - credits_used: Number of credits consumed.
                - credits_remaining: Remaining credit balance.
        """
        resp = self.client.post(
            f"{self.base_url}/v1/search",
            json={
                "smiles_list": smiles_list,
                "max_results": max_results,
                "substructure_search": substructure_search,
                "max_depth": max_depth,
                "include_chemistry_risk": include_chemistry_risk,
                "include_chemistry_risk_score": include_chemistry_risk_score,
            },
        )
        if not resp.is_success:
            print(f"Error {resp.status_code}: {resp.text}")
        resp.raise_for_status()
        return resp.json()

    def order(
        self,
        smiles: list[str],
        email: str,
        notes: str = "",
    ):
        """Submit an order for the given molecules.

        Args:
            smiles: List of SMILES strings to order.
            email: Contact email for the order.
            notes: Optional free-text notes.

        Returns:
            Dict with:
                - order_id: UUID string identifying the order.
                - molecule_count: Number of molecules in the order.
        """
        resp = self.client.post(
            f"{self.base_url}/v1/order",
            json={
                "smiles": smiles,
                "email": email,
                "notes": notes,
            },
        )
        if not resp.is_success:
            print(f"Error {resp.status_code}: {resp.text}")
        resp.raise_for_status()
        return resp.json()
