import hashlib
import os
import threading
from collections import defaultdict
from typing import Any, Dict, Optional

import openai
from dotenv import load_dotenv

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

load_dotenv()

# Load API keys from environment variables
DEFAULT_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
DEFAULT_OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
DEFAULT_VERCEL_API_KEY = os.getenv("AI_GATEWAY_API_KEY")
DEFAULT_PRIME_INTELLECT_BASE_URL = "https://api.pinference.ai/api/v1/"


class OpenAIClient(BaseLM):
    """
    LM Client for running models with the OpenAI API. Works with vLLM as well.

    PERF-003: Implements connection pooling via get_instance() classmethod.
    """

    # PERF-003: Class-level instance cache for connection pooling
    _instances: Dict[str, "OpenAIClient"] = {}
    _lock = threading.Lock()

    @classmethod
    def get_instance(
        cls,
        api_key: str | None = None,
        model_name: str | None = None,
        base_url: str | None = None,
        **kwargs,
    ) -> "OpenAIClient":
        """
        Get a cached client instance or create a new one.

        PERF-003: Connection pooling - reuses existing SDK instances.

        Args:
            api_key: API key for authentication
            model_name: Model name to use
            base_url: Base URL for API
            **kwargs: Additional arguments

        Returns:
            Cached or new OpenAIClient instance
        """
        # Create cache key from configuration
        key_parts = [
            api_key or "",
            model_name or "",
            base_url or "",
        ]
        cache_key = hashlib.md5(":".join(key_parts).encode()).hexdigest()

        with cls._lock:
            if cache_key not in cls._instances:
                cls._instances[cache_key] = cls(
                    api_key=api_key,
                    model_name=model_name,
                    base_url=base_url,
                    **kwargs
                )
            return cls._instances[cache_key]

    @classmethod
    def clear_instances(cls) -> None:
        """Clear the instance cache. Useful for testing."""
        with cls._lock:
            cls._instances.clear()

    def __init__(
        self,
        api_key: str | None = None,
        model_name: str | None = None,
        base_url: str | None = None,
        **kwargs,
    ):
        super().__init__(model_name=model_name, **kwargs)

        if api_key is None:
            if base_url == "https://api.openai.com/v1" or base_url is None:
                api_key = DEFAULT_OPENAI_API_KEY
            elif base_url == "https://openrouter.ai/api/v1":
                api_key = DEFAULT_OPENROUTER_API_KEY
            elif base_url == "https://ai-gateway.vercel.sh/v1":
                api_key = DEFAULT_VERCEL_API_KEY

        # For vLLM, set base_url to local vLLM server address.
        self.client = openai.OpenAI(api_key=api_key, base_url=base_url)
        self.async_client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.model_name = model_name

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

    def completion(self, prompt: str | list[dict[str, Any]], model: str | None = None) -> str:
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for OpenAI client.")

        extra_body = {}
        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
            extra_body["usage"] = {"include": True}

        response = self.client.chat.completions.create(
            model=model, messages=messages, extra_body=extra_body
        )
        self._track_cost(response, model)
        return response.choices[0].message.content

    async def acompletion(
        self, prompt: str | list[dict[str, Any]], model: str | None = None
    ) -> str:
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for OpenAI client.")

        extra_body = {}
        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
            extra_body["usage"] = {"include": True}

        response = await self.async_client.chat.completions.create(
            model=model, messages=messages, extra_body=extra_body
        )
        self._track_cost(response, model)
        return response.choices[0].message.content

    def _track_cost(self, response: openai.ChatCompletion, model: str):
        self.model_call_counts[model] += 1

        usage = getattr(response, "usage", None)
        if usage is None:
            raise ValueError("No usage data received. Tracking tokens not possible.")

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

        # Track last call for handler to read
        self.last_prompt_tokens = usage.prompt_tokens
        self.last_completion_tokens = usage.completion_tokens

    def get_usage_summary(self) -> UsageSummary:
        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:
        return ModelUsageSummary(
            total_calls=1,
            total_input_tokens=self.last_prompt_tokens,
            total_output_tokens=self.last_completion_tokens,
        )

    def completion_stream(
        self,
        prompt: str | list[dict[str, Any]],
        model: str | None = None,
        chunk_size: int = 1,  # OPT-004: Smaller default chunk size for faster first token
    ):
        """
        Stream completion response token by token.

        OPT-004: Optimized streaming with:
        - Reduced buffer sizes for earlier first token
        - Immediate yielding of chunks
        - Backpressure handling via generator pattern

        Args:
            prompt: Text prompt or message list
            model: Model name (optional, uses default)
            chunk_size: Minimum characters to buffer before yielding (default: 1)

        Yields:
            str: Response chunks as they arrive

        Example:
            for chunk in client.completion_stream("Hello"):
                print(chunk, end="", flush=True)
        """
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for OpenAI client.")

        extra_body = {}
        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
            extra_body["usage"] = {"include": True}

        # OPT-004: Stream response with small buffer
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            extra_body=extra_body,
        )

        buffer = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                buffer += content

                # OPT-004: Yield immediately when buffer exceeds chunk_size
                if len(buffer) >= chunk_size:
                    yield buffer
                    buffer = ""

        # Yield any remaining content
        if buffer:
            yield buffer

        # Track usage if available (may not be available in streaming)
        self.model_call_counts[model] += 1

    async def acompletion_stream(
        self,
        prompt: str | list[dict[str, Any]],
        model: str | None = None,
        chunk_size: int = 1,  # OPT-004: Smaller default chunk size
    ):
        """
        Async stream completion response token by token.

        OPT-004: Optimized async streaming with:
        - Reduced buffer sizes for earlier first token
        - Immediate yielding of chunks
        - Backpressure handling via async generator

        Args:
            prompt: Text prompt or message list
            model: Model name (optional, uses default)
            chunk_size: Minimum characters to buffer before yielding (default: 1)

        Yields:
            str: Response chunks as they arrive

        Example:
            async for chunk in client.acompletion_stream("Hello"):
                print(chunk, end="", flush=True)
        """
        if isinstance(prompt, str):
            messages = [{"role": "user", "content": prompt}]
        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
            messages = prompt
        else:
            raise ValueError(f"Invalid prompt type: {type(prompt)}")

        model = model or self.model_name
        if not model:
            raise ValueError("Model name is required for OpenAI client.")

        extra_body = {}
        if self.client.base_url == DEFAULT_PRIME_INTELLECT_BASE_URL:
            extra_body["usage"] = {"include": True}

        # OPT-004: Async stream response
        response = await self.async_client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            extra_body=extra_body,
        )

        buffer = ""
        async for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                buffer += content

                # OPT-004: Yield immediately when buffer exceeds chunk_size
                if len(buffer) >= chunk_size:
                    yield buffer
                    buffer = ""

        # Yield any remaining content
        if buffer:
            yield buffer

        # Track usage
        self.model_call_counts[model] += 1
