"""
Central AI Provider Service
Unified gateway for all AI completions — supports Groq and g4f backends.
All route files should import from here instead of calling g4f directly.

Usage:
    from services.ai_provider import ai_generate, ai_generate_stream, ai_generate_vision, ai_generate_async

    # Simple text generation
    result = ai_generate(messages, model='llama-3.3-70b')

    # Streaming
    for chunk in ai_generate_stream(messages, model='llama-3.3-70b'):
        print(chunk)

    # Vision
    result = ai_generate_vision(messages, images=[("data:image/png;base64,...", "photo.png")])

    # Async  (for quiz/qa routes)
    result = await ai_generate_async(messages, model='llama-3.3-70b')
"""

import logging
import time
import json
import base64
from dataclasses import dataclass, field
from typing import Optional, List, Tuple, Generator, Any

from services.ai_config import (
    AI_BACKEND, GROQ_API_KEY, GROQ_BASE_URL,
    GROQ_DEFAULT_MODEL, GROQ_VISION_MODEL,
    resolve_groq_model, resolve_groq_vision_model,
    get_groq_model_capabilities,
    is_groq_enabled, is_g4f_enabled,
    GROQ_MODEL_MAPPING,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Response wrapper
# ---------------------------------------------------------------------------

@dataclass
class AIResponse:
    """Normalized response from any AI backend."""
    content: str
    model: str = ''
    provider: str = ''          # 'groq' or 'g4f'
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    finish_reason: str = ''
    raw: Any = field(default=None, repr=False)


# ---------------------------------------------------------------------------
# Groq backend
# ---------------------------------------------------------------------------

_groq_client = None
_async_groq_client = None


def _get_groq_client():
    """Lazily create Groq client."""
    global _groq_client
    if _groq_client is None:
        from groq import Groq
        _groq_client = Groq(api_key=GROQ_API_KEY)
    return _groq_client


def _get_async_groq_client():
    """Lazily create async Groq client."""
    global _async_groq_client
    if _async_groq_client is None:
        from groq import AsyncGroq
        _async_groq_client = AsyncGroq(api_key=GROQ_API_KEY)
    return _async_groq_client


def _groq_generate(messages: list, model: str = None, stream: bool = False,
                   temperature: float = 0.7, max_tokens: int = None,
                   response_format: dict = None, **kwargs) -> AIResponse:
    """Generate text with Groq API (synchronous, non-streaming)."""
    client = _get_groq_client()
    groq_model = resolve_groq_model(model)

    create_kwargs = {
        'model': groq_model,
        'messages': messages,
        'temperature': max(temperature, 1e-8),  # Groq rejects exact 0
        'stream': False,
    }
    if max_tokens:
        create_kwargs['max_completion_tokens'] = max_tokens
    if response_format:
        create_kwargs['response_format'] = response_format

    completion = client.chat.completions.create(**create_kwargs)

    choice = completion.choices[0]
    usage = completion.usage or {}

    return AIResponse(
        content=choice.message.content or '',
        model=groq_model,
        provider='groq',
        input_tokens=getattr(usage, 'prompt_tokens', 0),
        output_tokens=getattr(usage, 'completion_tokens', 0),
        total_tokens=getattr(usage, 'total_tokens', 0),
        finish_reason=getattr(choice, 'finish_reason', ''),
        raw=completion,
    )


def _groq_generate_stream(messages: list, model: str = None,
                           temperature: float = 0.7, max_tokens: int = None,
                           **kwargs) -> Generator[str, None, None]:
    """Generate text with Groq API (streaming). Yields content strings."""
    client = _get_groq_client()
    groq_model = resolve_groq_model(model)

    create_kwargs = {
        'model': groq_model,
        'messages': messages,
        'temperature': max(temperature, 1e-8),
        'stream': True,
    }
    if max_tokens:
        create_kwargs['max_completion_tokens'] = max_tokens

    stream = client.chat.completions.create(**create_kwargs)

    for chunk in stream:
        if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content


def _groq_generate_vision(messages: list, images: list,
                            model: str = None, temperature: float = 0.7,
                            max_tokens: int = 1024, **kwargs) -> AIResponse:
    """Generate response for vision request via Groq.
    
    Args:
        messages: Conversation messages (system + user).
        images: List of (data_uri_or_url, filename) tuples.
        model: Model name (will be resolved to vision-capable model).
    """
    client = _get_groq_client()
    vision_model = resolve_groq_vision_model()

    # Build content array for the last user message with image blocks
    formatted_messages = []
    for msg in messages:
        if msg['role'] == 'user' and images:
            content_blocks = []
            # Add text
            if msg.get('content'):
                content_blocks.append({
                    'type': 'text',
                    'text': msg['content']
                })
            # Add images
            for img in images:
                if isinstance(img, (list, tuple)):
                    img_data = img[0]
                elif isinstance(img, str):
                    img_data = img
                else:
                    continue

                # Handle base64 data URIs and raw base64
                if hasattr(img_data, 'read'):
                    # File-like object — read and encode
                    raw_bytes = img_data.read()
                    img_data = f"data:image/jpeg;base64,{base64.b64encode(raw_bytes).decode()}"
                elif isinstance(img_data, bytes):
                    img_data = f"data:image/jpeg;base64,{base64.b64encode(img_data).decode()}"
                elif isinstance(img_data, str) and not img_data.startswith(('data:', 'http')):
                    img_data = f"data:image/jpeg;base64,{img_data}"

                content_blocks.append({
                    'type': 'image_url',
                    'image_url': {'url': img_data}
                })
            formatted_messages.append({'role': 'user', 'content': content_blocks})
            images = None  # Only attach images to first user message
        else:
            formatted_messages.append(msg)

    completion = client.chat.completions.create(
        model=vision_model,
        messages=formatted_messages,
        temperature=max(temperature, 1e-8),
        max_completion_tokens=max_tokens,
        stream=False,
    )

    choice = completion.choices[0]
    usage = completion.usage or {}

    return AIResponse(
        content=choice.message.content or '',
        model=vision_model,
        provider='groq',
        input_tokens=getattr(usage, 'prompt_tokens', 0),
        output_tokens=getattr(usage, 'completion_tokens', 0),
        total_tokens=getattr(usage, 'total_tokens', 0),
        finish_reason=getattr(choice, 'finish_reason', ''),
        raw=completion,
    )


async def _groq_generate_async(messages: list, model: str = None,
                                temperature: float = 0.7, max_tokens: int = None,
                                response_format: dict = None, **kwargs) -> AIResponse:
    """Generate text with Groq API (async)."""
    client = _get_async_groq_client()
    groq_model = resolve_groq_model(model)

    create_kwargs = {
        'model': groq_model,
        'messages': messages,
        'temperature': max(temperature, 1e-8),
        'stream': False,
    }
    if max_tokens:
        create_kwargs['max_completion_tokens'] = max_tokens
    if response_format:
        create_kwargs['response_format'] = response_format

    completion = await client.chat.completions.create(**create_kwargs)

    choice = completion.choices[0]
    usage = completion.usage or {}

    return AIResponse(
        content=choice.message.content or '',
        model=groq_model,
        provider='groq',
        input_tokens=getattr(usage, 'prompt_tokens', 0),
        output_tokens=getattr(usage, 'completion_tokens', 0),
        total_tokens=getattr(usage, 'total_tokens', 0),
        finish_reason=getattr(choice, 'finish_reason', ''),
        raw=completion,
    )


# ---------------------------------------------------------------------------
# g4f backend (wraps existing g4f calls)
# ---------------------------------------------------------------------------

def _g4f_generate(messages: list, model: str = None, provider: str = None,
                  stream: bool = False, images: list = None, **kwargs) -> AIResponse:
    """Generate text with g4f (synchronous, non-streaming)."""
    import g4f

    create_kwargs = {
        'model': model or 'gemini-search',
        'messages': messages,
        'stream': False,
    }
    if provider:
        create_kwargs['provider'] = provider
    if images:
        create_kwargs['images'] = images

    response = g4f.ChatCompletion.create(**create_kwargs)
    content = response if isinstance(response, str) else str(response)

    # Estimate tokens from word count (g4f doesn't provide usage)
    input_words = sum(len(m.get('content', '').split()) for m in messages if isinstance(m.get('content'), str))
    output_words = len(content.split()) if content else 0

    return AIResponse(
        content=content,
        model=model or 'gemini-search',
        provider='g4f',
        input_tokens=input_words,
        output_tokens=output_words,
        total_tokens=input_words + output_words,
    )


def _g4f_generate_stream(messages: list, model: str = None,
                          provider: str = None, **kwargs) -> Generator[str, None, None]:
    """Generate text with g4f (streaming). Yields content strings."""
    import g4f

    create_kwargs = {
        'model': model or 'command-r-plus-08-2024',
        'messages': messages,
        'stream': True,
    }
    if provider:
        create_kwargs['provider'] = provider or 'CohereForAI_C4AI_Command'

    response = g4f.ChatCompletion.create(**create_kwargs)
    for chunk in response:
        if isinstance(chunk, str):
            yield chunk
        elif isinstance(chunk, dict):
            if 'choices' in chunk and chunk['choices']:
                delta = chunk['choices'][0].get('delta', {})
                text = delta.get('content', '')
                if text:
                    yield text
            elif 'content' in chunk:
                yield chunk['content']


def _g4f_generate_vision(messages: list, images: list,
                          model: str = None, provider: str = None, **kwargs) -> AIResponse:
    """Generate response for vision request via g4f."""
    import g4f

    create_kwargs = {
        'model': model or '',
        'messages': messages,
        'stream': False,
        'images': images,
    }
    if provider:
        create_kwargs['provider'] = provider or 'Copilot'

    response = g4f.ChatCompletion.create(**create_kwargs)
    content = response if isinstance(response, str) else str(response)

    input_words = sum(len(m.get('content', '').split()) for m in messages if isinstance(m.get('content'), str))
    output_words = len(content.split()) if content else 0

    return AIResponse(
        content=content,
        model=model or 'vision-auto',
        provider='g4f',
        input_tokens=input_words + 1000,  # 1000 for image
        output_tokens=output_words,
        total_tokens=input_words + output_words + 1000,
    )


async def _g4f_generate_async(messages: list, model: str = None,
                               provider: str = None, **kwargs) -> AIResponse:
    """Generate text with g4f (async)."""
    import g4f

    create_kwargs = {
        'model': model or 'gemini-search',
        'messages': messages,
    }
    if provider:
        create_kwargs['provider'] = provider

    response = await g4f.ChatCompletion.create_async(**create_kwargs)
    content = response if isinstance(response, str) else str(response)

    input_words = sum(len(m.get('content', '').split()) for m in messages if isinstance(m.get('content'), str))
    output_words = len(content.split()) if content else 0

    return AIResponse(
        content=content,
        model=model or 'gemini-search',
        provider='g4f',
        input_tokens=input_words,
        output_tokens=output_words,
        total_tokens=input_words + output_words,
    )


# ---------------------------------------------------------------------------
# Public API — these are the functions all routes should call
# ---------------------------------------------------------------------------

def _select_backend(model: str = None, backend_override: str = None) -> str:
    """Decide which backend to use.
    Returns 'groq' or 'g4f'.
    """
    if backend_override:
        return backend_override

    # If a model has a DB-level backend_type, that takes precedence
    # (handled by the caller via get_model_config)

    if AI_BACKEND == 'groq' and is_groq_enabled():
        return 'groq'
    elif AI_BACKEND == 'g4f':
        return 'g4f'
    elif AI_BACKEND == 'auto':
        return 'groq' if is_groq_enabled() else 'g4f'
    else:
        return 'g4f'  # ultimate fallback


def ai_generate(messages: list, model: str = None, stream: bool = False,
                temperature: float = 0.7, max_tokens: int = None,
                response_format: dict = None,
                backend: str = None,
                g4f_provider: str = None,
                images: list = None,
                **kwargs) -> AIResponse:
    """
    Generate AI text response (synchronous, non-streaming).
    
    Args:
        messages: List of {"role": ..., "content": ...} dicts.
        model: Model name (will be mapped to backend-specific ID).
        stream: Must be False for this function.
        temperature: Sampling temperature.
        max_tokens: Max output tokens.
        response_format: e.g. {"type": "json_object"} for JSON mode.
        backend: Force 'groq' or 'g4f'. None = use global config.
        g4f_provider: g4f-specific provider string (ignored for Groq).
        images: For vision — list of (data_uri, filename) tuples.
        
    Returns:
        AIResponse with content, usage, model info.
    """
    chosen = _select_backend(model, backend)

    if chosen == 'groq':
        try:
            if images:
                return _groq_generate_vision(messages, images, model=model,
                                              temperature=temperature, max_tokens=max_tokens or 1024, **kwargs)
            return _groq_generate(messages, model=model, temperature=temperature,
                                  max_tokens=max_tokens, response_format=response_format, **kwargs)
        except Exception as e:
            logger.error(f"Groq generate failed: {e}")
            if AI_BACKEND == 'auto' and is_g4f_enabled():
                logger.info("Falling back to g4f...")
                return _g4f_generate(messages, model=model, provider=g4f_provider,
                                     images=images, **kwargs)
            raise
    else:
        return _g4f_generate(messages, model=model, provider=g4f_provider,
                             images=images, **kwargs)


def ai_generate_stream(messages: list, model: str = None,
                        temperature: float = 0.7, max_tokens: int = None,
                        backend: str = None,
                        g4f_provider: str = None,
                        **kwargs) -> Generator[str, None, None]:
    """
    Generate AI text response (streaming). Yields content chunk strings.
    
    Use inside an SSE generator to stream chunks to the client.
    """
    chosen = _select_backend(model, backend)

    if chosen == 'groq':
        try:
            yield from _groq_generate_stream(messages, model=model,
                                              temperature=temperature, max_tokens=max_tokens, **kwargs)
        except Exception as e:
            logger.error(f"Groq streaming failed: {e}")
            if AI_BACKEND == 'auto' and is_g4f_enabled():
                logger.info("Falling back to g4f streaming...")
                yield from _g4f_generate_stream(messages, model=model,
                                                 provider=g4f_provider, **kwargs)
            else:
                raise
    else:
        yield from _g4f_generate_stream(messages, model=model,
                                         provider=g4f_provider, **kwargs)


def ai_generate_vision(messages: list, images: list,
                        model: str = None, temperature: float = 0.7,
                        max_tokens: int = 1024,
                        backend: str = None,
                        g4f_provider: str = None,
                        **kwargs) -> AIResponse:
    """
    Generate AI response for vision/image analysis requests.
    
    Args:
        messages: Conversation messages.
        images: List of (data_uri_or_url, filename) or data URIs.
        model: Model name (auto-selects vision model for Groq).
    """
    chosen = _select_backend(model, backend)

    if chosen == 'groq':
        try:
            return _groq_generate_vision(messages, images, model=model,
                                          temperature=temperature, max_tokens=max_tokens, **kwargs)
        except Exception as e:
            logger.error(f"Groq vision failed: {e}")
            if AI_BACKEND == 'auto' and is_g4f_enabled():
                logger.info("Falling back to g4f vision...")
                return _g4f_generate_vision(messages, images, model=model,
                                             provider=g4f_provider, **kwargs)
            raise
    else:
        return _g4f_generate_vision(messages, images, model=model,
                                     provider=g4f_provider, **kwargs)


async def ai_generate_async(messages: list, model: str = None,
                             temperature: float = 0.7, max_tokens: int = None,
                             response_format: dict = None,
                             backend: str = None,
                             g4f_provider: str = None,
                             **kwargs) -> AIResponse:
    """
    Generate AI text response (async, for use in async quiz/qa generators).
    """
    chosen = _select_backend(model, backend)

    if chosen == 'groq':
        try:
            return await _groq_generate_async(messages, model=model,
                                               temperature=temperature, max_tokens=max_tokens,
                                               response_format=response_format, **kwargs)
        except Exception as e:
            logger.error(f"Groq async failed: {e}")
            if AI_BACKEND == 'auto' and is_g4f_enabled():
                logger.info("Falling back to g4f async...")
                return await _g4f_generate_async(messages, model=model,
                                                  provider=g4f_provider, **kwargs)
            raise
    else:
        return await _g4f_generate_async(messages, model=model,
                                          provider=g4f_provider, **kwargs)


def get_current_backend() -> str:
    """Return which backend is currently active."""
    return _select_backend()


def get_provider_info() -> dict:
    """Return info about current AI provider configuration."""
    return {
        'backend': AI_BACKEND,
        'active_backend': _select_backend(),
        'groq_enabled': is_groq_enabled(),
        'g4f_enabled': is_g4f_enabled(),
        'groq_default_model': GROQ_DEFAULT_MODEL,
        'groq_vision_model': GROQ_VISION_MODEL,
    }
