src/proofframe/providers.py

"""Media generation providers."""

from __future__ import annotations

import html
import mimetypes
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol
from urllib.parse import urlparse
from urllib.request import urlopen

from .checksum import sha256_hex
from .config import ConfigurationError, Settings
from .models import Campaign, GeneratedMedia


class MediaProvider(Protocol):
    provider: str
    model: str

    def generate(self, campaign: Campaign, count: int = 3) -> list[GeneratedMedia]: ...


def build_campaign_prompt(campaign: Campaign, index: int) -> str:
    return (
        f"Create a {campaign.tone} campaign image for {campaign.audience}. "
        f"Brief: {campaign.brief}. Variant {index}."
    )


def build_local_genblaze_svg(prompt: str, model: str, step_id: str) -> tuple[bytes, str]:
    seed = sha256_hex(f"{step_id}:{model}:{prompt}".encode("utf-8"))[:8]
    colors = ["#2F6B5F", "#C94C4C", "#F2C14E", "#3D5A80", "#8E5CF0"]
    primary = colors[int(seed[:2], 16) % len(colors)]
    secondary = colors[int(seed[2:4], 16) % len(colors)]
    safe_prompt = html.escape(prompt[:120])
    safe_model = html.escape(model)
    svg = f"""<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="720" viewBox="0 0 1280 720">
  <rect width="1280" height="720" fill="#14171A"/>
  <rect x="72" y="72" width="1136" height="576" rx="18" fill="#F7F4EA"/>
  <rect x="120" y="126" width="380" height="26" fill="{primary}"/>
  <circle cx="1056" cy="188" r="88" fill="{secondary}" opacity="0.9"/>
  <text x="120" y="238" font-family="Helvetica, Arial, sans-serif" font-size="58" fill="#161616">ProofFrame</text>
  <text x="124" y="314" font-family="Helvetica, Arial, sans-serif" font-size="28" fill="#333333">{safe_prompt}</text>
  <text x="124" y="550" font-family="Helvetica, Arial, sans-serif" font-size="22" fill="#555555">Genblaze local provider | model {safe_model} | seed {seed}</text>
</svg>
"""
    return svg.encode("utf-8"), seed


@dataclass(frozen=True)
class MockMediaProvider:
    """Deterministic media provider for tests and credential-free demos."""

    model: str = "mock-svg-v1"
    provider: str = "mock"

    def generate(self, campaign: Campaign, count: int = 3) -> list[GeneratedMedia]:
        return [self._generate_one(campaign, index) for index in range(1, count + 1)]

    def _generate_one(self, campaign: Campaign, index: int) -> GeneratedMedia:
        prompt = build_campaign_prompt(campaign, index)
        seed = sha256_hex(f"{campaign.id}:{prompt}".encode("utf-8"))[:8]
        colors = ["#F2C14E", "#4F7CAC", "#6FAE75", "#C45BAA", "#E4572E"]
        primary = colors[int(seed[:2], 16) % len(colors)]
        secondary = colors[int(seed[2:4], 16) % len(colors)]
        title = html.escape(campaign.title[:42])
        brief = html.escape(campaign.brief[:96])
        svg = f"""<svg xmlns="http://www.w3.org/2000/svg" width="1280" height="720" viewBox="0 0 1280 720">
  <rect width="1280" height="720" fill="#101216"/>
  <rect x="80" y="76" width="1120" height="568" rx="28" fill="#F8F4EA"/>
  <circle cx="1080" cy="180" r="92" fill="{primary}" opacity="0.92"/>
  <rect x="140" y="142" width="420" height="28" fill="{secondary}"/>
  <text x="140" y="252" font-family="Georgia, serif" font-size="72" fill="#151515">{title}</text>
  <text x="144" y="326" font-family="Helvetica, Arial, sans-serif" font-size="28" fill="#333">{brief}</text>
  <text x="144" y="548" font-family="Helvetica, Arial, sans-serif" font-size="22" fill="#555">ProofFrame mock asset | seed {seed} | variant {index}</text>
</svg>
"""
        return GeneratedMedia(
            prompt=prompt,
            provider=self.provider,
            model=self.model,
            filename=f"{campaign.id}-variant-{index}.svg",
            content_type="image/svg+xml",
            data=svg.encode("utf-8"),
            generation_metadata={"mock_seed": seed, "variant": str(index)},
        )


@dataclass(frozen=True)
class GenblazeMediaProvider:
    """Genblaze image generation adapter using the official Pipeline API."""

    api_key: str
    image_model: str
    genblaze_provider: str = "gmicloud"
    base_url: str = ""
    aspect_ratio: str = "16:9"
    timeout_seconds: int = 180
    b2_sink_enabled: bool = False
    b2_bucket: str = ""
    b2_key_id: str = ""
    b2_application_key: str = ""
    b2_region: str = ""
    b2_public_base_url: str = ""
    provider: str = "genblaze"

    @property
    def model(self) -> str:
        return self.image_model

    @classmethod
    def from_settings(cls, settings: Settings) -> "GenblazeMediaProvider":
        settings.require_genblaze()
        return cls(
            api_key=settings.genblaze_provider_key(),
            image_model=settings.genblaze_image_model,
            genblaze_provider=settings.genblaze_provider,
            base_url=settings.genblaze_base_url,
            aspect_ratio=settings.genblaze_aspect_ratio,
            timeout_seconds=settings.genblaze_timeout_seconds,
            b2_sink_enabled=settings.storage_backend == "b2",
            b2_bucket=settings.b2_bucket,
            b2_key_id=settings.b2_key_id,
            b2_application_key=settings.b2_application_key,
            b2_region=settings.b2_region_for_backblaze(),
            b2_public_base_url=settings.b2_public_base_url,
        )

    def generate(self, campaign: Campaign, count: int = 3) -> list[GeneratedMedia]:
        try:
            from genblaze_core import Modality, Pipeline  # type: ignore[import-not-found]
        except ModuleNotFoundError as exc:
            raise ConfigurationError(
                "Genblaze generation requires the official Genblaze core package. "
                "Install with `pip install -e '.[integrations]'` and configure a live provider."
            ) from exc

        generated: list[GeneratedMedia] = []
        for index in range(1, count + 1):
            prompt = build_campaign_prompt(campaign, index)
            provider, provider_suffix = self._build_genblaze_provider()
            sink, storage_backend = self._build_genblaze_b2_sink(campaign.id)
            try:
                result = (
                    Pipeline("proofframe-image", project_id=campaign.id)
                    .step(
                        provider,
                        model=self.image_model,
                        prompt=prompt,
                        modality=Modality.IMAGE,
                        aspect_ratio=self.aspect_ratio,
                    )
                    .run(
                        sink=sink,
                        timeout=self.timeout_seconds,
                        max_retries=1,
                        raise_on_failure=True,
                    )
                )
                run = getattr(result, "run", None)
                manifest = getattr(result, "manifest", None)
                if run is None or manifest is None:
                    run, manifest = result
                step = run.steps[0]
                if not step.assets:
                    raise ConfigurationError("Genblaze generation completed without an image asset.")

                asset = step.assets[0]
                try:
                    data, content_type, filename = self._fetch_asset(
                        asset.url,
                        campaign.id,
                        index,
                        storage_backend=storage_backend,
                    )
                except Exception:
                    if self.genblaze_provider != "local":
                        raise
                    data, _seed = build_local_genblaze_svg(prompt, self.image_model, step.step_id)
                    content_type = "image/svg+xml"
                    filename = f"{campaign.id}-genblaze-{index}.svg"
                manifest_uri = str(getattr(manifest, "manifest_uri", "") or "")
                generated.append(
                    GeneratedMedia(
                        prompt=prompt,
                        provider=f"{self.provider}/{provider_suffix}",
                        model=self.image_model,
                        filename=filename,
                        content_type=content_type or asset.media_type or "image/png",
                        data=data,
                        generation_metadata={
                            "genblaze_b2_sink": "enabled" if self.b2_sink_enabled else "disabled",
                            "genblaze_run_id": str(run.run_id),
                            "genblaze_manifest_hash": str(manifest.canonical_hash),
                            "genblaze_manifest_uri_present": str(bool(manifest_uri)),
                            "genblaze_manifest_verified": str(manifest.verify()),
                            "genblaze_step_status": str(step.status),
                            "genblaze_asset_host": urlparse(asset.url).netloc or "local",
                            "genblaze_asset_sha256": str(asset.sha256 or ""),
                        },
                    )
                )
            except ConfigurationError:
                raise
            except Exception as exc:
                raise ConfigurationError(f"Genblaze generation failed: {exc}") from exc
            finally:
                self._close_if_possible(storage_backend)
                self._close_if_possible(provider)
        return generated

    def _build_genblaze_provider(self) -> tuple[Any, str]:
        if self.genblaze_provider == "gmicloud":
            try:
                from genblaze_gmicloud import GMICloudImageProvider  # type: ignore[import-not-found]
            except ModuleNotFoundError as exc:
                raise ConfigurationError(
                    "GENBLAZE_PROVIDER=gmicloud requires genblaze-gmicloud. "
                    "Install with `pip install -e '.[integrations]'`."
                ) from exc
            return (
                GMICloudImageProvider(
                    api_key=self.api_key,
                    base_url=self.base_url or None,
                    http_timeout=float(self.timeout_seconds),
                ),
                "gmicloud-image",
            )
        if self.genblaze_provider == "openai":
            try:
                from genblaze_openai import DalleProvider  # type: ignore[import-not-found]
            except ModuleNotFoundError as exc:
                raise ConfigurationError(
                    "GENBLAZE_PROVIDER=openai requires genblaze-openai. "
                    "Install with `pip install -e '.[integrations]'`."
                ) from exc
            return (
                DalleProvider(
                    api_key=self.api_key,
                    http_timeout=float(self.timeout_seconds),
                ),
                "openai-image",
            )
        if self.genblaze_provider == "local":
            try:
                from genblaze_core.models.asset import Asset  # type: ignore[import-not-found]
                from genblaze_core.providers.base import SyncProvider  # type: ignore[import-not-found]
            except ModuleNotFoundError as exc:
                raise ConfigurationError(
                    "GENBLAZE_PROVIDER=local requires genblaze-core. "
                    "Install with `pip install -e '.[integrations]'`."
                ) from exc

            class ProofFrameLocalImageProvider(SyncProvider):  # type: ignore[misc, valid-type]
                name = "proofframe-local"

                def __init__(self) -> None:
                    super().__init__()
                    self.output_dir = Path(tempfile.gettempdir()) / "proofframe-genblaze-local"
                    self.output_dir.mkdir(parents=True, exist_ok=True)

                def generate(self, step: Any, config: Any | None = None) -> Any:
                    prompt = str(step.prompt or "")
                    data, seed = build_local_genblaze_svg(
                        prompt,
                        str(step.model),
                        step.step_id,
                    )
                    output_path = self.output_dir / f"{step.step_id}.svg"
                    output_path.write_bytes(data)
                    asset = Asset(
                        url=output_path.as_uri(),
                        media_type="image/svg+xml",
                        sha256=sha256_hex(data),
                        size_bytes=len(data),
                        width=1280,
                        height=720,
                        metadata={"provider_mode": "credential_free_local", "seed": seed},
                    )
                    step.assets.append(asset)
                    step.provider_payload = {
                        "provider_mode": "credential_free_local",
                        "model": str(step.model),
                    }
                    return step

            return (ProofFrameLocalImageProvider(), "local-image")
        raise ConfigurationError(
            "Unsupported Genblaze provider: "
            f"{self.genblaze_provider}. Use GENBLAZE_PROVIDER=gmicloud, openai, or local."
        )

    def _build_genblaze_b2_sink(
        self,
        campaign_id: str,
        *,
        preflight: bool = True,
    ) -> tuple[Any | None, Any | None]:
        if not self.b2_sink_enabled:
            return None, None
        missing = [
            name
            for name, value in {
                "B2_BUCKET": self.b2_bucket,
                "B2_KEY_ID": self.b2_key_id,
                "B2_APPLICATION_KEY or B2_APP_KEY": self.b2_application_key,
                "B2_REGION or Backblaze S3 endpoint": self.b2_region,
            }.items()
            if not value
        ]
        if missing:
            raise ConfigurationError(
                "Genblaze B2 sink was requested but required environment variables are missing: "
                + ", ".join(missing)
            )
        try:
            from genblaze_core import KeyStrategy, ObjectStorageSink  # type: ignore[import-not-found]
            from genblaze_s3 import S3StorageBackend  # type: ignore[import-not-found]
        except ModuleNotFoundError as exc:
            raise ConfigurationError(
                "Genblaze B2 sink requires the official genblaze-s3 package. "
                "Install with `pip install -e '.[integrations]'`."
            ) from exc

        try:
            backend = S3StorageBackend.for_backblaze(
                self.b2_bucket,
                region=self.b2_region,
                key_id=self.b2_key_id,
                app_key=self.b2_application_key,
                public_url_base=self.b2_public_base_url or None,
                preflight=preflight,
            )
            read_backend = S3StorageBackend.for_backblaze(
                self.b2_bucket,
                region=self.b2_region,
                key_id=self.b2_key_id,
                app_key=self.b2_application_key,
                public_url_base=self.b2_public_base_url or None,
                preflight=False,
            )
            sink = ObjectStorageSink(
                backend,
                prefix=f"campaigns/{campaign_id}/genblaze",
                key_strategy=KeyStrategy.HIERARCHICAL,
            )
        except Exception as exc:
            raise ConfigurationError(f"Genblaze B2 sink configuration failed: {exc}") from exc
        return sink, read_backend

    @staticmethod
    def _close_if_possible(resource: Any | None) -> None:
        close = getattr(resource, "close", None)
        if callable(close):
            close()

    @staticmethod
    def _fetch_asset(
        asset_url: str,
        campaign_id: str,
        index: int,
        *,
        storage_backend: Any | None = None,
    ) -> tuple[bytes, str, str]:
        parsed = urlparse(asset_url)
        if storage_backend is not None:
            key_from_url = getattr(storage_backend, "key_from_url", None)
            get = getattr(storage_backend, "get", None)
            if callable(key_from_url) and callable(get):
                key = key_from_url(asset_url)
                if key:
                    data = get(key)
                    content_type = mimetypes.guess_type(key)[0] or "image/png"
                    extension = mimetypes.guess_extension(content_type) or ".png"
                    return data, content_type, f"{campaign_id}-genblaze-{index}{extension}"
        if parsed.scheme in {"http", "https"}:
            with urlopen(asset_url, timeout=60) as response:
                data = response.read()
                content_type = response.headers.get_content_type()
        else:
            path = Path(parsed.path if parsed.scheme == "file" else asset_url)
            data = path.read_bytes()
            content_type = mimetypes.guess_type(path.name)[0] or "image/png"

        extension = mimetypes.guess_extension(content_type) or ".png"
        return data, content_type, f"{campaign_id}-genblaze-{index}{extension}"


def create_media_provider(settings: Settings) -> MediaProvider:
    if settings.generation_backend == "mock":
        return MockMediaProvider()
    if settings.generation_backend == "genblaze":
        return GenblazeMediaProvider.from_settings(settings)
    raise ConfigurationError(f"Unknown generation backend: {settings.generation_backend}")