"""
MiniMax Helper Library
A comprehensive helper class for interacting with MiniMax API.
Provides similar functionality to the OpenAI helper but for MiniMax endpoints.
"""

import os
import asyncio
import json
import time
import logging
import base64
from datetime import datetime, timedelta
from typing import List, Dict, Optional, Union, Callable, Any, AsyncGenerator, Generator
from dataclasses import dataclass, field
from pathlib import Path
import hashlib
from .exceptions import InsufficientQuotaError

try:
    import requests
    from requests.adapters import HTTPAdapter
    from urllib3.util.retry import Retry
except ImportError:
    raise ImportError("Please install requests: pip install requests")


@dataclass
class MiniMaxConfig:
    """Configuration class for MiniMax helper."""
    api_key: Optional[str] = None
    api_secret: Optional[str] = None
    base_url: str = "https://api.minimax.io/v1"
    chat_model: str = "MiniMax-M2"
    embedding_model: str = "embo-01"
    image_model: str = "image-01"
    tts_model: str = "speech-01"
    stt_model: str = "speech-01"
    vision_model: str = "MiniMax-M2"
    group_id: Optional[str] = None
    max_retries: int = 3
    timeout: float = 60.0
    enable_caching: bool = True
    cache_ttl: int = 3600  # 1 hour
    enable_cost_tracking: bool = True
    log_level: str = "INFO"

    def __post_init__(self):
        if not self.api_key:
            self.api_key = os.getenv("MINIMAX_API_KEY")
        if not self.api_key:
            raise EnvironmentError("MINIMAX_API_KEY not set in environment or config.")

        if not self.api_secret:
            self.api_secret = os.getenv("MINIMAX_API_SECRET")

        if not self.group_id:
            self.group_id = os.getenv("MINIMAX_GROUP_ID")

        # Env overrides
        self.chat_model = os.getenv("MINIMAX_CHAT_MODEL", self.chat_model)
        self.embedding_model = os.getenv("MINIMAX_EMBED_MODEL", self.embedding_model)
        self.image_model = os.getenv("MINIMAX_IMAGE_MODEL", self.image_model)


@dataclass
class CacheEntry:
    """Cache entry with TTL support."""
    data: Any
    created_at: datetime
    ttl: int

    def is_expired(self) -> bool:
        return datetime.now() > self.created_at + timedelta(seconds=self.ttl)


@dataclass
class UsageStats:
    """Usage statistics tracking for MiniMax API."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    requests: int = 0
    estimated_cost: float = 0.0

    # Central pricing table: approximate, easy to extend.
    # Values are USD per 1K tokens. Update as needed.
    PRICE = {
        # Chat/completion (input/output) - approximate pricing
        "MiniMax-M2": {"input": 0.0, "output": 0.0},
        "MiniMax-M1": {"input": 0.0, "output": 0.0},
        "abab6.5s-chat": {"input": 0.0, "output": 0.0},
        # Embeddings (input only)
        "embo-01": {"input": 0.0001, "output": 0.0},
        # Image generation
        "image-01": {"input": 0.0, "output": 0.0},
    }

    @staticmethod
    def _model_family(model: str) -> Optional[str]:
        """Map a specific model name to a pricing family key."""
        if not model:
            return None
        model_lower = model.lower()
        for family in UsageStats.PRICE.keys():
            if family.lower() in model_lower:
                return family
        return None

    def add_usage(self, prompt_tokens: int = 0, completion_tokens: int = 0, model: str = ""):
        self.prompt_tokens += prompt_tokens
        self.completion_tokens += completion_tokens
        self.total_tokens += prompt_tokens + completion_tokens
        self.requests += 1

        family = self._model_family(model)
        if not family:
            return  # unknown model; skip cost estimate

        rates = self.PRICE[family]
        self.estimated_cost += (prompt_tokens / 1000.0) * rates.get("input", 0.0)
        self.estimated_cost += (completion_tokens / 1000.0) * rates.get("output", 0.0)


class MiniMaxHelper:
    """Enhanced MiniMax helper with advanced features."""

    def __init__(self, config: Optional[MiniMaxConfig] = None):
        self.config = config or MiniMaxConfig()
        self._setup_logging()
        self._setup_session()
        self._setup_cache()
        self._setup_usage_tracking()

    def _setup_logging(self):
        logging.basicConfig(
            level=getattr(logging, self.config.log_level.upper(), logging.INFO),
            format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        )
        self.logger = logging.getLogger(self.__class__.__name__)

    def _setup_session(self):
        """Setup HTTP session with retry strategy."""
        self.session = requests.Session()
        retry_strategy = Retry(
            total=self.config.max_retries,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

    def _setup_cache(self):
        self._cache: Dict[str, CacheEntry] = {}

    def _setup_usage_tracking(self):
        self.usage_stats = UsageStats()

    # ---------------- HTTP helpers ----------------
    def _get_headers(self, content_type: str = "application/json") -> Dict[str, str]:
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": content_type,
        }
        return headers

    def _build_url(self, endpoint: str) -> str:
        base = self.config.base_url.rstrip("/")
        endpoint = endpoint.lstrip("/")
        url = f"{base}/{endpoint}"
        if self.config.group_id:
            url = f"{url}?GroupId={self.config.group_id}"
        return url

    # ---------------- Cache helpers ----------------
    def _cache_key(self, *args, **kwargs) -> str:
        key_data = {"args": args, "kwargs": kwargs}
        key_str = json.dumps(key_data, sort_keys=True, default=str)
        return hashlib.md5(key_str.encode()).hexdigest()

    def _get_from_cache(self, key: str) -> Optional[Any]:
        if not self.config.enable_caching:
            return None
        entry = self._cache.get(key)
        if entry and not entry.is_expired():
            self.logger.debug(f"Cache hit for key: {key}")
            return entry.data
        if entry:
            self._cache.pop(key, None)
        return None

    def _set_cache(self, key: str, data: Any, ttl: Optional[int] = None):
        if not self.config.enable_caching:
            return
        ttl = ttl or self.config.cache_ttl
        self._cache[key] = CacheEntry(data=data, created_at=datetime.now(), ttl=ttl)
        self.logger.debug(f"Cache set for key: {key}")

    # ---------------- Retry helper ----------------
    def _retry_with_backoff(self, func: Callable, *args, **kwargs):
        max_retries = self.config.max_retries
        base_delay = 1.0
        for attempt in range(max_retries + 1):
            try:
                return func(*args, **kwargs)
            except AttributeError as e:
                # Handle 'NoneType' object has no attribute 'request' or broken sessions
                if "'NoneType' object has no attribute 'request'" in str(e):
                    self.logger.warning(f"MiniMax session internal state is broken (attempt {attempt + 1}), re-initializing...")
                    self._setup_session()
                    if attempt == max_retries:
                        raise
                    continue
                raise
            except requests.exceptions.RequestException as e:
                if attempt == max_retries:
                    raise
                delay = base_delay * (2 ** attempt)
                self.logger.warning(f"Request error, retrying in {delay:.1f}s (attempt {attempt + 1}): {e}")
                time.sleep(delay)

    # ---------------- Chat ----------------
    def chat(
        self,
        messages: List[Dict[str, Any]],
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        stream: bool = False,
        tools: Optional[List[Dict]] = None,
        tool_choice: Optional[Union[str, Dict]] = None,
        stop: Optional[List[str]] = None,
        top_p: Optional[float] = None,
        **kwargs,
    ) -> Union[str, Dict, Generator[str, None, None]]:
        """Chat completion with function/tool calling support."""
        mdl = model or self.config.chat_model
        from .tracing_helper import trace_span
        with trace_span("minimax_chat", {"model": mdl, "stream": stream}):
            if not (0 <= temperature <= 2):
                raise ValueError("Temperature must be between 0 and 2")

            request_params = {
                "model": mdl,
                "messages": messages,
                "temperature": temperature,
                "stream": stream,
                **kwargs,
            }
            if max_tokens is not None:
                request_params["max_tokens"] = max_tokens
            if tools:
                request_params["tools"] = tools
            if tool_choice:
                request_params["tool_choice"] = tool_choice
            if stop:
                request_params["stop"] = stop
            if top_p is not None:
                request_params["top_p"] = top_p

            cache_key = None
            if not stream and self.config.enable_caching:
                cache_key = self._cache_key("chat", **request_params)
                cached = self._get_from_cache(cache_key)
                if cached is not None:
                    return cached

            def _make_request():
                url = self._build_url("chat/completions")

                headers = self._get_headers()

                resp = self.session.post(
                    url,
                    json=request_params,
                    headers=headers,
                    timeout=self.config.timeout,
                    stream=stream,
                )
                resp.raise_for_status()

                if stream:
                    def stream_generator():
                        chunks: List[str] = []
                        for line in resp.iter_lines():
                            if line:
                                line = line.decode('utf-8')
                                if line.startswith('data: '):
                                    data = line[6:]
                                    if data == '[DONE]':
                                        break
                                    try:
                                        chunk_data = json.loads(data)
                                        delta = (chunk_data.get("choices") or [{}])[0].get("delta", {})
                                        content = delta.get("content") or delta.get("reasoning_content") or ""
                                        if content:
                                            chunks.append(content)
                                            yield content
                                    except json.JSONDecodeError:
                                        continue
                        if cache_key:
                            self._set_cache(cache_key, "".join(chunks))
                    return stream_generator()
                else:
                    data = resp.json()
                    
                    # Check for API errors in base_resp FIRST
                    base_resp = data.get("base_resp", {})
                    if base_resp.get("status_code", 0) != 0:
                        status_msg = base_resp.get('status_msg', '')
                        error_msg = f"MiniMax API Error: {status_msg} (Code: {base_resp.get('status_code')})"
                        
                        # Check for quota/credit errors
                        if "quota" in status_msg.lower() or "credit" in status_msg.lower() or "insufficient" in status_msg.lower():
                            raise InsufficientQuotaError(error_msg)
                            
                        self.logger.error(error_msg)
                        # If there's an error, we might not have choices
                        if not data.get("choices"):
                            return error_msg

                    choice = (data.get("choices") or [{}])[0]
                    message = choice.get("message", {})

                    if message.get("tool_calls") or message.get("function_call"):
                        out = {
                            "content": message.get("content") or message.get("reasoning_content") or "",
                            "tool_calls": message.get("tool_calls"),
                            "function_call": message.get("function_call"),
                            "finish_reason": choice.get("finish_reason"),
                        }
                    else:
                        out = message.get("content") or message.get("reasoning_content") or ""

                    # Track usage
                    usage = data.get("usage", {})
                    tokens_used = usage.get("total_tokens", 0)
                    
                    if not out and tokens_used > 0:
                        self.logger.warning(f"MiniMax response content is empty but tokens were consumed ({tokens_used}). Raw message: {message}")

                    if usage and self.config.enable_cost_tracking:
                        self.usage_stats.add_usage(
                            prompt_tokens=usage.get("prompt_tokens", 0),
                            completion_tokens=usage.get("completion_tokens", 0),
                            model=mdl,
                        )

                    if cache_key:
                        self._set_cache(cache_key, out)
                    return out

            return self._retry_with_backoff(_make_request)

    async def achat(
        self,
        messages: List[Dict[str, Any]],
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None,
        **kwargs,
    ) -> Union[str, Dict]:
        """Async chat completion."""
        mdl = model or self.config.chat_model
        from .tracing_helper import trace_span
        async with trace_span("minimax_achat", {"model": mdl}):
            request_params = {
                "model": mdl,
                "messages": messages,
                "temperature": temperature,
                **kwargs,
            }
            if max_tokens is not None:
                request_params["max_tokens"] = max_tokens

            cache_key = None
            if self.config.enable_caching:
                cache_key = self._cache_key("achat", **request_params)
                cached = self._get_from_cache(cache_key)
                if cached is not None:
                    return cached

            url = self._build_url("chat/completions")
            headers = self._get_headers()

            loop = asyncio.get_event_loop()
            resp = await loop.run_in_executor(
                None,
                lambda: requests.post(url, json=request_params, headers=headers, timeout=self.config.timeout)
            )
            resp.raise_for_status()
            data = resp.json()

            # Check for API errors in base_resp FIRST
            base_resp = data.get("base_resp", {})
            if base_resp.get("status_code", 0) != 0:
                status_msg = base_resp.get('status_msg', '')
                error_msg = f"MiniMax API Error (Async): {status_msg} (Code: {base_resp.get('status_code')})"
                
                # Check for quota/credit errors
                if "quota" in status_msg.lower() or "credit" in status_msg.lower() or "insufficient" in status_msg.lower():
                    raise InsufficientQuotaError(error_msg)
                    
                self.logger.error(error_msg)
                if not data.get("choices"):
                    return error_msg

            choice = (data.get("choices") or [{}])[0]
            message = choice.get("message", {})

            if message.get("tool_calls"):
                out = {
                    "content": message.get("content") or message.get("reasoning_content") or "",
                    "tool_calls": message.get("tool_calls"),
                    "finish_reason": choice.get("finish_reason"),
                }
            else:
                out = message.get("content") or message.get("reasoning_content") or ""

            usage = data.get("usage", {})
            if usage and self.config.enable_cost_tracking:
                self.usage_stats.add_usage(
                    prompt_tokens=usage.get("prompt_tokens", 0),
                    completion_tokens=usage.get("completion_tokens", 0),
                    model=mdl,
                )

            if cache_key:
                self._set_cache(cache_key, out)
            return out

    # ---------------- Vision in chat ----------------
    def chat_with_vision(
        self,
        text: str,
        images: List[Union[str, Path]],
        model: Optional[str] = None,
        **kwargs
    ) -> str:
        """Chat with vision capabilities."""
        model = model or self.config.vision_model
        content: List[Dict[str, Any]] = [{"type": "text", "text": text}]

        for img in images:
            if isinstance(img, (str, Path)):
                s = str(img)
                if s.startswith(("http://", "https://")):
                    content.append({"type": "image_url", "image_url": {"url": s}})
                else:
                    with open(s, "rb") as f:
                        b64 = base64.b64encode(f.read()).decode()
                    content.append({"type": "image_url", "image_url": {"url": f"data:image/*;base64,{b64}"}})

        messages = [{"role": "user", "content": content}]
        return self.chat(messages, model=model, **kwargs)

    # ---------------- Embeddings ----------------
    def get_embedding(
        self,
        text: str,
        model: Optional[str] = None,
        dimensions: Optional[int] = None,
    ) -> List[float]:
        """Get embedding for a single text."""
        mdl = model or self.config.embedding_model
        from .tracing_helper import trace_span
        with trace_span("minimax_embedding", {"model": mdl}):
            cache_key = self._cache_key("embedding", text=text, model=mdl, dimensions=dimensions)
            cached = self._get_from_cache(cache_key)
            if cached is not None:
                return cached

            def _make_request():
                request_data = {
                    "model": mdl,
                    "texts": [text],
                }
                if dimensions:
                    request_data["dimensions"] = dimensions

                url = self._build_url("embeddings")
                headers = self._get_headers()

                resp = self.session.post(url, json=request_data, headers=headers, timeout=self.config.timeout)
                resp.raise_for_status()
                data = resp.json()

                embedding = data.get("data", [{}])[0].get("embedding", [])

                usage = data.get("usage", {})
                if usage and self.config.enable_cost_tracking:
                    self.usage_stats.add_usage(
                        prompt_tokens=usage.get("prompt_tokens", 0),
                        model=mdl,
                    )

                self._set_cache(cache_key, embedding)
                return embedding

            return self._retry_with_backoff(_make_request)

    def get_embeddings(
        self,
        texts: List[str],
        model: Optional[str] = None,
        dimensions: Optional[int] = None,
        batch_size: int = 100,
    ) -> List[List[float]]:
        """Get embeddings for multiple texts."""
        mdl = model or self.config.embedding_model
        from .tracing_helper import trace_span
        with trace_span("minimax_embeddings", {"model": mdl, "count": len(texts)}):
            all_embeddings: List[List[float]] = []
            total_batches = (len(texts) + batch_size - 1) // batch_size

            for idx in range(0, len(texts), batch_size):
                batch = texts[idx:idx + batch_size]

                def _make_request():
                    request_data = {
                        "model": mdl,
                        "texts": batch,
                    }
                    if dimensions:
                        request_data["dimensions"] = dimensions

                    url = self._build_url("embeddings")
                    headers = self._get_headers()

                    resp = self.session.post(url, json=request_data, headers=headers, timeout=self.config.timeout)
                    resp.raise_for_status()
                    data = resp.json()

                    embeddings = [item.get("embedding", []) for item in data.get("data", [])]

                    usage = data.get("usage", {})
                    if usage and self.config.enable_cost_tracking:
                        self.usage_stats.add_usage(
                            prompt_tokens=usage.get("prompt_tokens", 0),
                            model=mdl,
                        )

                    return embeddings

                embeddings = self._retry_with_backoff(_make_request)
                all_embeddings.extend(embeddings)
                self.logger.info(f"Processed batch {idx // batch_size + 1}/{total_batches}")

            return all_embeddings

    # ---------------- Images ----------------
    def generate_image(
        self,
        prompt: str,
        model: Optional[str] = None,
        size: str = "1024x1024",
        quality: str = "standard",
        n: int = 1,
        style: str = "natural",
    ) -> List[str]:
        """Generate images with MiniMax image generation API."""
        mdl = model or self.config.image_model
        from .tracing_helper import trace_span
        with trace_span("minimax_image_gen", {"model": mdl, "n": n, "size": size}):
            def _make_request():
                request_data = {
                    "model": mdl,
                    "prompt": prompt,
                    "num_images": n,
                    "size": size,
                    "quality": quality,
                    "style": style,
                }

                url = self._build_url("text2image")
                headers = self._get_headers()

                resp = self.session.post(url, json=request_data, headers=headers, timeout=self.config.timeout * 5)
                resp.raise_for_status()
                data = resp.json()

                self.usage_stats.requests += 1
                images = data.get("data", [])
                return [img.get("url", "") for img in images]

            return self._retry_with_backoff(_make_request)

    # ---------------- Audio ----------------
    def text_to_speech(
        self,
        text: str,
        voice: str = "male-qn-qingse",
        model: Optional[str] = None,
        speed: float = 1.0,
        vol: float = 1.0,
        pitch: int = 0,
        format: str = "mp3",
        bitrate: int = 128000,
        sample_rate: int = 32000,
    ) -> bytes:
        """Convert text to speech."""
        mdl = model or self.config.tts_model
        from .tracing_helper import trace_span
        with trace_span("minimax_tts", {"model": mdl, "voice": voice}):
            def _make_request():
                request_data = {
                    "model": mdl,
                    "text": text,
                    "voice_setting": {
                        "voice_id": voice,
                        "speed": speed,
                        "vol": vol,
                        "pitch": pitch,
                    },
                    "audio_setting": {
                        "format": format,
                        "bitrate": bitrate,
                        "sample_rate": sample_rate,
                    },
                }

                url = self._build_url("text/speech")
                headers = self._get_headers()

                resp = self.session.post(url, json=request_data, headers=headers, timeout=self.config.timeout)
                resp.raise_for_status()

                self.usage_stats.requests += 1
                return resp.content

            return self._retry_with_backoff(_make_request)

    def speech_to_text(
        self,
        audio_file: Union[str, Path],
        model: Optional[str] = None,
        language: Optional[str] = None,
        prompt: Optional[str] = None,
    ) -> str:
        """Convert speech to text."""
        mdl = model or self.config.stt_model
        from .tracing_helper import trace_span
        with trace_span("minimax_stt", {"model": mdl}):
            def _make_request():
                url = self._build_url("text/speech2text")

                # Build headers - for multipart, we need to handle differently
                headers = {}
                if self.config.api_key:
                    headers["Authorization"] = f"Bearer {self.config.api_key}"

                with open(audio_file, "rb") as f:
                    files = {"file": (os.path.basename(audio_file), f, "audio/mpeg")}
                    data = {"model": mdl}
                    if language:
                        data["language"] = language
                    if prompt:
                        data["prompt"] = prompt

                    resp = self.session.post(url, files=files, data=data, headers=headers, timeout=self.config.timeout)

                resp.raise_for_status()
                result = resp.json()

                self.usage_stats.requests += 1
                return result.get("text", "")

            return self._retry_with_backoff(_make_request)

    # ---------------- Token estimation (approximate) ----------------
    def estimate_tokens(self, text: str) -> int:
        """Estimate token count (approximate)."""
        # Rough estimate: 1 token ≈ 4 characters for English, 2 for Chinese
        # This is a simple approximation
        chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        other_chars = len(text) - chinese_chars
        return (chinese_chars // 2) + (other_chars // 4)

    def count_tokens(self, text: str, model: Optional[str] = None) -> int:
        """Count tokens for a given text."""
        model = model or self.config.chat_model

        cache_key = self._cache_key("tokens", text=text, model=model)
        cached = self._get_from_cache(cache_key)
        if cached is not None:
            return cached

        tokens = self.estimate_tokens(text)
        self._set_cache(cache_key, tokens)
        return tokens

    def manage_context_length(
        self,
        messages: List[Dict[str, str]],
        max_tokens: int = 4000,
        preserve_system: bool = True,
    ) -> List[Dict[str, str]]:
        """Trim oldest messages to fit max_tokens. Keeps system message if requested."""
        def msg_tokens(m: Dict[str, str]) -> int:
            return self.estimate_tokens(m.get("content", ""))

        total = sum(msg_tokens(m) for m in messages)
        if total <= max_tokens:
            return messages

        kept: List[Dict[str, str]] = []
        work = messages[:]

        if preserve_system and work and work[0].get("role") == "system":
            kept.append(work.pop(0))
            total = sum(msg_tokens(m) for m in kept + work)

        while work and total > max_tokens:
            removed = work.pop(0)
            total = sum(msg_tokens(m) for m in kept + work)

        result = kept + work
        self.logger.info(f"Context management: kept {len(result)} messages (~{total} tokens)")
        return result

    # ---------------- Models ----------------
    def list_models(self) -> List[str]:
        """List available models."""
        cache_key = "models_list"
        cached = self._get_from_cache(cache_key)
        if cached is not None:
            return cached

        def _make_request():
            url = self._build_url("models")
            headers = self._get_headers()

            resp = self.session.get(url, headers=headers, timeout=self.config.timeout)
            resp.raise_for_status()
            data = resp.json()

            models = [m.get("id", "") for m in data.get("data", [])]
            self._set_cache(cache_key, models, ttl=86400)
            return models

        return self._retry_with_backoff(_make_request)

    # ---------------- Introspection & utilities ----------------
    def get_usage_stats(self) -> Dict[str, Any]:
        return {
            "prompt_tokens": self.usage_stats.prompt_tokens,
            "completion_tokens": self.usage_stats.completion_tokens,
            "total_tokens": self.usage_stats.total_tokens,
            "requests": self.usage_stats.requests,
            "estimated_cost": round(self.usage_stats.estimated_cost, 6),
            "cache_size": len(self._cache),
        }

    def clear_cache(self):
        """Clear the cache."""
        self._cache.clear()
        self.logger.info("Cache cleared")

    def reset_usage_stats(self):
        """Reset usage statistics."""
        self.usage_stats = UsageStats()
        self.logger.info("Usage statistics reset")

    def batch_process(
        self,
        items: List[Any],
        processor: Callable[[Any], Any],
        batch_size: int = 10,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> List[Any]:
        """Process items in batches."""
        results: List[Any] = []
        total_batches = (len(items) + batch_size - 1) // batch_size

        for i in range(0, len(items), batch_size):
            batch = items[i:i + batch_size]
            batch_results: List[Any] = []

            for item in batch:
                try:
                    batch_results.append(processor(item))
                except Exception as e:
                    self.logger.error(f"Error processing item {item}: {e}")
                    batch_results.append(None)

            results.extend(batch_results)

            if progress_callback:
                progress_callback(i // batch_size + 1, total_batches)
            self.logger.info(f"Processed batch {i // batch_size + 1}/{total_batches}")

        return results

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        stats = self.get_usage_stats()
        self.logger.info(f"Session complete. Usage: {stats}")


# Type alias for streaming generator (imported from typing)


# Convenience functions for message creation
def create_system_message(content: str) -> Dict[str, str]:
    """Create a system message."""
    return {"role": "system", "content": content}


def create_user_message(content: str) -> Dict[str, str]:
    """Create a user message."""
    return {"role": "user", "content": content}


def create_assistant_message(content: str, tool_calls: Optional[list] = None) -> Dict[str, Any]:
    """Create an assistant message."""
    message = {"role": "assistant", "content": content}
    if tool_calls:
        message["tool_calls"] = tool_calls
    return message


def create_tool_message(tool_call_id: str, content: str) -> Dict[str, str]:
    """Create a tool result message."""
    return {
        "role": "tool",
        "tool_call_id": tool_call_id,
        "content": content,
    }
