import json
from collections.abc import Generator

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 search_stream(
        self,
        smiles: str,
        max_results: int = 100,
        substructure_search: bool = False,
        include_chemistry_risk: bool = False,
        include_chemistry_risk_score: bool = False,
    ) -> Generator[dict]:
        """Stream search results for a single molecule via Server-Sent Events.

        Yields dicts with 'status' and 'message' keys as the search progresses.
        Status lifecycle: starting → synthesis → rescoring → complete.
        The final event (status='complete') includes a 'results' list with the
        same fields as the batch search endpoint.

        Args:
            smiles: Single SMILES string to search.
            max_results: Maximum results (default 100).
            substructure_search: Match substructures instead of similarity.
            include_chemistry_risk: Include chemistry_risk field (low/medium/high).
            include_chemistry_risk_score: Include raw probability score (0-1).

        Yields:
            Dict with 'status' and 'message'. Final event also contains 'results'.
        """
        with self.client.stream(
            "POST",
            f"{self.base_url}/v1/search/stream",
            json={
                "smiles": smiles,
                "max_results": max_results,
                "substructure_search": substructure_search,
                "include_chemistry_risk": include_chemistry_risk,
                "include_chemistry_risk_score": include_chemistry_risk_score,
            },
        ) as resp:
            if not resp.is_success:
                raise httpx.HTTPStatusError(
                    f"Error {resp.status_code}", request=resp.request, response=resp
                )
            buffer = ""
            for chunk in resp.iter_text():
                buffer += chunk
                while "\n\n" in buffer:
                    event, buffer = buffer.split("\n\n", 1)
                    for line in event.strip().split("\n"):
                        if line.startswith("data: "):
                            yield json.loads(line[6:])


    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()
