import os
import time
from specific_ai.common.inference import Inference
from specific_ai.common.tracing import Tracing


class SpecificAIClient:
    def __init__(
        self,
        url: str = None,
        trace: bool = False,
        triton_url: str = None,
        *,
        base_url: str | None = None,
        api_key: str | None = None,
        api_key_header: str = "Authorization",
        api_key_prefix: str | None = "Bearer",
        timeout_s: float = 30.0,
        session_cookie: str | None = None,
        session_cookie_name: str = "session",
    ):
        """
        Initialize the SpecificAI SDK client.

        This client supports two modes:
        1) Inference mode (existing behavior): provide `url` (gateway) or `triton_url` (direct).
        2) Platform mode (new): provide `base_url` + `api_key` to manage tasks/datasets/trainings.

        Args:
            url: SpecificAI gateway URL for inference (required for tracing).
            trace: Whether to send traces back to the platform (requires `url`).
            triton_url: Triton inference URL (alternative inference mode).
            base_url: Base URL of the SpecificAI backend API for platform actions.
            api_key: API key / bearer token for platform actions.
            api_key_header: Header name used to send the API key.
            api_key_prefix: Optional prefix (e.g. "Bearer") for the API key header.
            timeout_s: Default HTTP timeout for platform actions.
            session_cookie: Optional session cookie value (for cookie-auth deployments).
            session_cookie_name: Cookie name for `session_cookie`.

        Raises:
            ValueError: If neither inference nor platform configuration is provided.
        """
        self.inference = None
        self.tracing = None
        self.should_trace = False

        # Platform client (tasks/datasets/trainings/models)
        self.platform = None
        self.tasks = None
        self.assets = None
        self.setup = None
        self.trainings = None
        self.models = None

        # Allow environment-based configuration for a simpler UX:
        #   SpecificAIClient(api_key="...") with SPECIFIC_AI_BASE_URL set.
        if base_url is None:
            base_url = os.getenv("SPECIFIC_AI_BASE_URL")
        if api_key is None:
            api_key = os.getenv("SPECIFIC_AI_API_KEY")

        if base_url:
            from specific_ai.platform.base_client import BaseClient
            from specific_ai.platform.client import SpecificAIPlatformClient

            http = BaseClient(
                base_url=base_url,
                api_key=api_key,
                api_key_header=api_key_header,
                api_key_prefix=api_key_prefix,
                timeout_s=timeout_s,
                session_cookie=session_cookie,
                session_cookie_name=session_cookie_name,
            )
            self.platform = SpecificAIPlatformClient(http)
            # Convenience aliases
            self.tasks = self.platform.tasks
            self.assets = self.platform.assets
            # Backwards-compatible alias: model setup is a task concern.
            self.setup = self.platform.tasks
            self.trainings = self.platform.trainings
            self.models = self.platform.models

        # Inference client (existing behavior)
        if url is not None or triton_url is not None:
            if url is None and trace:
                raise ValueError("'url' is required when 'trace' is True")

            self.inference = Inference(url, triton_url)
            self.tracing = Tracing(url) if url else None
            self.should_trace = trace

        if self.inference is None and self.platform is None:
            raise ValueError(
                "Provide either inference config (`url` or `triton_url`) or platform config (`base_url`)."
            )

    def create(self, message: str, task_name: str, project_name: str = "default"):
        """
        Run inference on a deployed SpecificAI model (inference mode only).

        Args:
            message: Prompt/message to send.
            task_name: Task name in the SpecificAI platform.
            project_name: Project name in the SpecificAI platform.

        Returns:
            LMResponse: Parsed model response (type depends on task).

        Raises:
            ValueError: If the client was initialized without inference config.
            Exception: If inference fails.
        """
        if self.inference is None:
            raise ValueError(
                "Inference is not configured. Initialize with `url` or `triton_url`."
            )

        response = None
        raw_logits = None
        inference_error = None

        try:
            start_time = time.time()
            responses, raw_logits = self.inference.infer(
                message, task_name, project_name
            )
            response = responses[0]
            return response
        except Exception as e:
            inference_error = str(e)
            raise Exception("Error inferring on specific.ai model") from e
        finally:
            if self.should_trace:
                self.tracing.collect(
                    response=response,
                    raw_logits=raw_logits,
                    inference_error=inference_error,
                    modelname="specific.ai",
                    prompt=message,
                    usecase_name=task_name,
                    usecase_group=project_name,
                    response_time=time.time() - start_time,
                    is_from_specific_ai_model=True,
                    logprobs=None,
                )
