"""
Cohere Client - Enterprise AI for Production

Provides LLM completion via Cohere API.
Cohere specializes in enterprise-grade AI with strong RAG and embedding capabilities.
"""

import os
from collections import defaultdict
from typing import Any, Optional

from groknroll.clients.base_lm import BaseLM
from groknroll.core.types import ModelUsageSummary, UsageSummary


class CohereClient(BaseLM):
    """
    Cohere Language Model Client

    Cohere provides enterprise-grade AI models:
    - Command R+ (128k context, best for RAG)
    - Command R (best value)
    - Command (production use)
    - Embed v3 (embeddings)
    - Rerank (search reranking)

    Features:
    - Chat completion
    - Streaming responses
    - RAG-optimized models
    - Grounded generation with citations
    - Multi-lingual support
    - Token usage tracking

    Example:
        >>> from groknroll.clients.cohere import CohereClient
        >>> client = CohereClient(model_name="command-r-plus")
        >>> response = await client.acompletion("Summarize this document")
        >>> print(response)
    """

    # Cohere models
    MODELS = [
        "command-r-plus",  # Latest R+, 128k context, best for RAG
        "command-r",  # Command R, best value
        "command",  # Standard Command model
        "command-light",  # Faster, lighter
        "embed-english-v3.0",  # English embeddings
        "embed-multilingual-v3.0",  # Multilingual embeddings
    ]

    def __init__(
        self,
        api_key: Optional[str] = None,
        model_name: str = "command-r-plus",
        base_url: Optional[str] = None,
        **kwargs,
    ):
        """
        Initialize Cohere client

        Args:
            api_key: Cohere API key (reads from COHERE_API_KEY env if not provided)
            model_name: Model to use (default: command-r-plus)
            base_url: Optional base URL override
            **kwargs: Additional Cohere-specific options
        """
        super().__init__(model_name=model_name, **kwargs)

        self.api_key = api_key or os.getenv("COHERE_API_KEY")
        if not self.api_key:
            raise ValueError(
                "Cohere API key not found. Set COHERE_API_KEY environment variable or pass api_key parameter."
            )

        self.model_name = model_name
        self.base_url = base_url

        # Initialize Cohere SDK
        try:
            import cohere
        except ImportError:
            raise ImportError(
                "cohere package not found. Install with: pip install cohere"
            )

        self.client = cohere.Client(api_key=self.api_key)
        self.async_client = cohere.AsyncClient(api_key=self.api_key)

        # Per-model usage tracking
        self.model_call_counts: dict[str, int] = defaultdict(int)
        self.model_input_tokens: dict[str, int] = defaultdict(int)
        self.model_output_tokens: dict[str, int] = defaultdict(int)
        self.model_total_tokens: dict[str, int] = defaultdict(int)

        # Last call tracking
        self.last_prompt_tokens = 0
        self.last_completion_tokens = 0

    def _messages_to_cohere_format(self, messages: list[dict[str, Any]]) -> tuple[str, Optional[str]]:
        """Convert OpenAI-style messages to Cohere format (message, preamble)"""
        preamble = None
        message = ""

        for msg in messages:
            role = msg.get("role")
            content = msg.get("content", "")

            if role == "system":
                preamble = content
            elif role == "user":
                message = content
            elif role == "assistant":
                # For chat history, we'd need to use chat_history parameter
                # For now, just append to message
                message += f"\nAssistant: {content}\n"

        return message, preamble

    def completion(
        self,
        prompt: str | list[dict[str, Any]],
        model: Optional[str] = None,
    ) -> str:
        """
        Synchronous completion

        Args:
            prompt: User prompt (string or messages list)
            model: Model to use (defaults to self.model_name)

        Returns:
            Generated text string
        """
        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Cohere client.")

        # Convert to Cohere format
        if isinstance(prompt, str):
            message = prompt
            preamble = None
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            message, preamble = self._messages_to_cohere_format(prompt)
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        response = self.client.chat(
            model=model,
            message=message,
            preamble=preamble,
        )

        self._track_cost(response, model)
        return response.text

    async def acompletion(
        self,
        prompt: str | list[dict[str, Any]],
        model: Optional[str] = None,
    ) -> str:
        """
        Asynchronous completion

        Args:
            prompt: User prompt (string or messages list)
            model: Model to use (defaults to self.model_name)

        Returns:
            Generated text string
        """
        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for Cohere client.")

        # Convert to Cohere format
        if isinstance(prompt, str):
            message = prompt
            preamble = None
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            message, preamble = self._messages_to_cohere_format(prompt)
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        response = await self.async_client.chat(
            model=model,
            message=message,
            preamble=preamble,
        )

        self._track_cost(response, model)
        return response.text

    def _track_cost(self, response, model: str):
        """Track token usage from API response"""
        self.model_call_counts[model] += 1

        # Cohere uses meta.tokens for usage
        if hasattr(response, "meta") and response.meta and hasattr(response.meta, "tokens"):
            tokens = response.meta.tokens
            prompt_tokens = tokens.input_tokens if hasattr(tokens, "input_tokens") else 0
            completion_tokens = tokens.output_tokens if hasattr(tokens, "output_tokens") else 0

            self.model_input_tokens[model] += prompt_tokens
            self.model_output_tokens[model] += completion_tokens
            self.model_total_tokens[model] += prompt_tokens + completion_tokens

            # Track last call for handler to read
            self.last_prompt_tokens = prompt_tokens
            self.last_completion_tokens = completion_tokens
        else:
            # No usage data available
            self.last_prompt_tokens = 0
            self.last_completion_tokens = 0

    def get_usage_summary(self) -> UsageSummary:
        """Get usage statistics across all models"""
        model_summaries = {}
        for model in self.model_call_counts:
            model_summaries[model] = ModelUsageSummary(
                total_calls=self.model_call_counts[model],
                total_input_tokens=self.model_input_tokens[model],
                total_output_tokens=self.model_output_tokens[model],
            )
        return UsageSummary(model_usage_summaries=model_summaries)

    def get_last_usage(self) -> ModelUsageSummary:
        """Get usage statistics from last completion"""
        return ModelUsageSummary(
            total_calls=1,
            total_input_tokens=self.last_prompt_tokens,
            total_output_tokens=self.last_completion_tokens,
        )
