persona-community-5/pkg/mediagen/adapters/laozhang.go

package adapters

import (
	"context"
	"encoding/base64"
	"fmt"

	"git.threesix.ai/jordan/persona-community-5/pkg/laozhang"
	"git.threesix.ai/jordan/persona-community-5/pkg/mediagen"
)

// LaoZhangProvider adapts pkg/laozhang to the mediagen interfaces.
type LaoZhangProvider struct {
	client *laozhang.Client
}

// NewLaoZhangProvider creates a new LaoZhang provider adapter.
func NewLaoZhangProvider(client *laozhang.Client) *LaoZhangProvider {
	return &LaoZhangProvider{client: client}
}

// Name implements mediagen.ImageGenerator and mediagen.VideoGenerator.
func (p *LaoZhangProvider) Name() string {
	return "laozhang"
}

// Health implements mediagen.ImageGenerator and mediagen.VideoGenerator.
func (p *LaoZhangProvider) Health(ctx context.Context) error {
	if err := p.client.Health(ctx); err != nil {
		return mediagen.NewProviderError("laozhang", "Health", err)
	}
	return nil
}

// GenerateImage implements mediagen.ImageGenerator.
// Supports reference images for identity consistency via Gemini native API.
// Seed support varies by model - it will be sent but may be ignored.
func (p *LaoZhangProvider) GenerateImage(ctx context.Context, req mediagen.ImageRequest) (*mediagen.ImageResponse, error) {
	// Map unified request to LaoZhang-specific request
	lzReq := laozhang.ImageRequest{
		Prompt:      req.Prompt,
		Model:       req.Model,
		Size:        req.Size,
		AspectRatio: req.AspectRatio,
		N:           req.Count,
		Seed:        req.Seed,
	}

	// Add reference image for identity consistency if provided
	if len(req.ReferenceImage) > 0 {
		mimeType := req.ReferenceMime
		if mimeType == "" {
			mimeType = "image/png"
		}
		lzReq.ReferenceImages = []laozhang.ReferenceImage{
			{
				Data:     req.ReferenceImage,
				MimeType: mimeType,
			},
		}
	}

	// Set defaults if not provided
	if lzReq.N == 0 {
		lzReq.N = 1
	}

	// Call LaoZhang client
	lzResp, err := p.client.GenerateImage(ctx, lzReq)
	if err != nil {
		return nil, mediagen.NewProviderError("laozhang", "GenerateImage", err)
	}

	// Convert response to unified format
	images := make([]mediagen.Image, 0, len(lzResp.Data))
	for _, img := range lzResp.Data {
		var data []byte
		var mimeType string
		var url string

		if img.B64JSON != "" {
			// Decode base64 to bytes
			decoded, err := base64.StdEncoding.DecodeString(img.B64JSON)
			if err != nil {
				return nil, fmt.Errorf("decode base64 image: %w", err)
			}
			data = decoded
			mimeType = "image/png" // LaoZhang returns PNG by default
		} else if img.URL != "" {
			url = img.URL
			mimeType = "image/png"
		}

		images = append(images, mediagen.Image{
			Data:     data,
			MimeType: mimeType,
			URL:      url,
		})
	}

	return &mediagen.ImageResponse{
		Images: images,
		Seed:   req.Seed, // Echo back seed (LaoZhang API doesn't return it)
	}, nil
}

// GenerateVideo implements mediagen.VideoGenerator.
func (p *LaoZhangProvider) GenerateVideo(ctx context.Context, req mediagen.VideoRequest) (*mediagen.VideoResponse, error) {
	// Determine model based on aspect ratio
	// Veo 3.1: default is portrait (9:16), add "-landscape" suffix for 16:9
	model := req.Model
	if model == "" {
		model = "veo-3.1"
	}
	if req.AspectRatio == "16:9" {
		model = model + "-landscape"
	}

	// Map unified request to LaoZhang-specific request
	lzReq := laozhang.VideoRequest{
		Prompt: req.Prompt,
		Model:  model,
		N:      req.Count,
	}

	// Set defaults if not provided
	if lzReq.N == 0 {
		lzReq.N = 1
	}

	// Convert reference images to base64 data URLs if provided
	if len(req.ReferenceImages) > 0 {
		lzReq.ReferenceImages = make([]string, 0, len(req.ReferenceImages))
		for _, img := range req.ReferenceImages {
			if len(img.Data) > 0 {
				// Convert bytes to data URL
				b64 := base64.StdEncoding.EncodeToString(img.Data)
				mimeType := img.MimeType
				if mimeType == "" {
					mimeType = "image/png"
				}
				dataURL := fmt.Sprintf("data:%s;base64,%s", mimeType, b64)
				lzReq.ReferenceImages = append(lzReq.ReferenceImages, dataURL)
			} else if img.URL != "" {
				lzReq.ReferenceImages = append(lzReq.ReferenceImages, img.URL)
			}
		}
	}

	// Call LaoZhang client
	lzResp, err := p.client.GenerateVideo(ctx, lzReq)
	if err != nil {
		return nil, mediagen.NewProviderError("laozhang", "GenerateVideo", err)
	}

	// Convert response to unified format - download video bytes from URLs
	videos := make([]mediagen.Video, 0, len(lzResp.Data))
	for _, vid := range lzResp.Data {
		if vid.URL == "" {
			continue
		}

		// Download video from URL
		data, err := downloadVideo(ctx, vid.URL)
		if err != nil {
			return nil, fmt.Errorf("download video from %s: %w", vid.URL, err)
		}

		videos = append(videos, mediagen.Video{
			Data:     data,
			URL:      vid.URL,
			MimeType: "video/mp4",
		})
	}

	if len(videos) == 0 {
		return nil, fmt.Errorf("no videos returned from LaoZhang")
	}

	return &mediagen.VideoResponse{
		Videos: videos,
	}, nil
}

// Compile-time interface check
var (
	_ mediagen.ImageGenerator = (*LaoZhangProvider)(nil)
	_ mediagen.VideoGenerator = (*LaoZhangProvider)(nil)
)