from __future__ import annotations

from typing import Any

from nees.client.transport import RemoteNEESTransport
from nees.orchestrator import Orchestrator
from nees.types import ProcessResponse


class NEES:
    def __init__(self, mode: str = "local", llm: str = "default", **remote_config: Any) -> None:
        if mode not in {"local", "remote"}:
            raise ValueError("mode must be 'local' or 'remote'")

        self._mode = mode
        self._llm = llm
        self._orchestrator = Orchestrator(llm=llm) if mode == "local" else None
        self._remote = (
            None
            if mode == "local"
            else RemoteNEESTransport(
                api_key=remote_config.get("api_key"),
                base_url=remote_config.get("base_url", "https://nees-gp.onrender.com"),
                timeout=remote_config.get("timeout", 30),
                max_retries=remote_config.get("max_retries", 2),
                verify_ssl=remote_config.get("verify_ssl", True),
            )
        )

    @property
    def mode(self) -> str:
        return self._mode

    def process(self, request: dict, enable_trace: bool = False) -> ProcessResponse:
        if self._mode == "local":
            assert self._orchestrator is not None
            return self._orchestrator.run(request, enable_trace=enable_trace)

        assert self._remote is not None
        return self._remote.process(request, enable_trace=enable_trace)

    def close(self) -> None:
        if self._remote is not None:
            self._remote.close()

    def __enter__(self) -> "NEES":
        return self

    def __exit__(self, *_: Any) -> None:
        self.close()
