from pathlib import Path

from harbor.llms.base import BaseLLM


class Chat:
    def __init__(self, model: BaseLLM):
        self._model = model
        self._messages = []
        self._cumulative_input_tokens = 0
        self._cumulative_output_tokens = 0
        self._cumulative_cache_tokens = 0
        self._cumulative_cost = 0.0
        self._all_logprobs: list[list[float]] = []

    @property
    def total_input_tokens(self) -> int:
        return self._cumulative_input_tokens

    @property
    def total_output_tokens(self) -> int:
        return self._cumulative_output_tokens

    @property
    def total_cache_tokens(self) -> int:
        return self._cumulative_cache_tokens

    @property
    def total_cost(self) -> float:
        return self._cumulative_cost

    @property
    def messages(self) -> list:
        return self._messages

    def get_all_logprobs(self) -> list[list[float]]:
        """Get all logprobs collected from all chat interactions.

        Returns:
            List of log probabilities for each assistant message.
            Each element is a list of logprobs for tokens in that message.
        """
        return self._all_logprobs

    async def chat(
        self,
        prompt: str,
        logging_path: Path | None = None,
        **kwargs,
    ) -> str:
        response = await self._model.call(
            prompt=prompt,
            message_history=self._messages,
            logging_path=logging_path,
            **kwargs,
        )

        # Get token usage and cost from the LLM response
        usage = self._model.get_last_usage()
        if usage is not None:
            self._cumulative_input_tokens += usage.prompt_tokens
            self._cumulative_output_tokens += usage.completion_tokens
            self._cumulative_cache_tokens += usage.cache_tokens
            self._cumulative_cost += usage.cost_usd

        # Collect logprobs if available
        logprobs = self._model.get_last_logprobs()
        if logprobs:
            self._all_logprobs.append(logprobs)

        self._messages.extend(
            [
                {"role": "user", "content": prompt},
                {"role": "assistant", "content": response},
            ]
        )
        return response
