SkillBridge AI — Test Prep Academy Platform

"""
AI provider abstraction. The rest of the app calls `get_provider().send_message(...)`
without knowing whether Groq or Anthropic is on the other end.

To swap providers: change AI_PROVIDER in .env. To add a new one: add a class here.

This is the single most important architectural decision in the backend. ~35
service modules call through this interface and none of them know which LLM
answers. Vision-model selection and token accounting are handled inside the
abstraction so calling code stays provider-agnostic.
"""

import base64
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path

from app.core.config import settings


@dataclass
class AttachedImage:
    """An image to send to the model. Either a local path or already-base64."""
    path: Path
    content_type: str  # e.g. "image/png", "image/jpeg"


@dataclass
class AIResponse:
    content: str
    input_tokens: int
    output_tokens: int


class AIProvider(ABC):
    """Contract every provider implementation must satisfy."""

    @abstractmethod
    def send_message(
        self,
        system_prompt: str,
        messages: list[dict],
        max_tokens: int = 800,
        images: list[AttachedImage] | None = None,
    ) -> AIResponse:
        """
        Send a conversation. `images` are attached to the LAST user message.
        Provider may switch to a vision model when images are present.
        """
        ...


def _read_image_as_data_url(image: AttachedImage) -> str:
    with image.path.open("rb") as f:
        encoded = base64.b64encode(f.read()).decode("ascii")
    return f"data:{image.content_type};base64,{encoded}"


class GroqProvider(AIProvider):
    """Groq + Llama. Switches to a vision model when images are present."""

    # Groq's vision-capable model. Llama-4 Scout is multimodal.
    _VISION_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"

    def __init__(self) -> None:
        from groq import Groq

        if not settings.groq_api_key:
            raise RuntimeError(
                "GROQ_API_KEY is not set. Add it to .env or switch AI_PROVIDER."
            )
        self._client = Groq(api_key=settings.groq_api_key)
        self._text_model = settings.groq_model

    def send_message(
        self,
        system_prompt: str,
        messages: list[dict],
        max_tokens: int = 800,
        images: list[AttachedImage] | None = None,
    ) -> AIResponse:
        # The branch that makes vision transparent to callers: if images are
        # present, build a multipart message and switch to the vision model.
        # Otherwise, plain text path on the default model.
        if images:
            full_messages = self._build_vision_messages(
                system_prompt, messages, images
            )
            model = self._VISION_MODEL
        else:
            full_messages = [
                {"role": "system", "content": system_prompt},
                *messages,
            ]
            model = self._text_model

        response = self._client.chat.completions.create(
            model=model,
            messages=full_messages,
            max_tokens=max_tokens,
            temperature=0.7,
        )

        usage = response.usage
        return AIResponse(
            content=response.choices[0].message.content or "",
            input_tokens=usage.prompt_tokens if usage else 0,
            output_tokens=usage.completion_tokens if usage else 0,
        )

    def _build_vision_messages(
        self,
        system_prompt: str,
        messages: list[dict],
        images: list[AttachedImage],
    ) -> list[dict]:
        """Convert the last user message to multipart with image_url blocks."""
        out: list[dict] = [{"role": "system", "content": system_prompt}]

        # Append all but the last verbatim
        for m in messages[:-1]:
            out.append(m)

        last = messages[-1]
        text = last.get("content") or ""
        content_parts = [{"type": "text", "text": text}]
        for img in images:
            content_parts.append(
                {
                    "type": "image_url",
                    "image_url": {"url": _read_image_as_data_url(img)},
                }
            )
        out.append({"role": last["role"], "content": content_parts})
        return out


class AnthropicProvider(AIProvider):
    """Anthropic Claude. Native multimodal."""

    def __init__(self) -> None:
        from anthropic import Anthropic

        if not settings.anthropic_api_key:
            raise RuntimeError(
                "ANTHROPIC_API_KEY is not set. Add it to .env or switch AI_PROVIDER."
            )
        self._client = Anthropic(api_key=settings.anthropic_api_key)
        self._model = settings.anthropic_model

    def send_message(
        self,
        system_prompt: str,
        messages: list[dict],
        max_tokens: int = 800,
        images: list[AttachedImage] | None = None,
    ) -> AIResponse:
        if images and messages:
            # Attach images to last user message in Anthropic's block format
            transformed = list(messages[:-1])
            last = messages[-1]
            blocks: list[dict] = [{"type": "text", "text": last.get("content") or ""}]
            for img in images:
                with img.path.open("rb") as f:
                    data = base64.b64encode(f.read()).decode("ascii")
                blocks.append(
                    {
                        "type": "image",
                        "source": {
                            "type": "base64",
                            "media_type": img.content_type,
                            "data": data,
                        },
                    }
                )
            transformed.append({"role": last["role"], "content": blocks})
            messages_to_send = transformed
        else:
            messages_to_send = messages

        response = self._client.messages.create(
            model=self._model,
            system=system_prompt,
            messages=messages_to_send,
            max_tokens=max_tokens,
        )

        text_parts = [
            block.text for block in response.content if hasattr(block, "text")
        ]
        return AIResponse(
            content="".join(text_parts),
            input_tokens=response.usage.input_tokens,
            output_tokens=response.usage.output_tokens,
        )


_provider_instance: AIProvider | None = None


def get_provider() -> AIProvider:
    """Return the configured provider, instantiated once and cached."""
    global _provider_instance
    if _provider_instance is not None:
        return _provider_instance

    name = settings.ai_provider.lower()
    if name == "groq":
        _provider_instance = GroqProvider()
    elif name == "anthropic":
        _provider_instance = AnthropicProvider()
    else:
        raise RuntimeError(
            f"Unknown AI_PROVIDER '{settings.ai_provider}'. "
            f"Use 'groq' or 'anthropic'."
        )

    return _provider_instance