package chatter

import (
	"context"
	"encoding/json"
	"fmt"
	"reflect"

	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/option"
)

type GeminiAdapter struct {
	client *genai.Client
	model  string
}

func NewGeminiAdapter(ctx context.Context, apiKey string, modelName string) (*GeminiAdapter, error) {
	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
	if err != nil {
		return nil, fmt.Errorf("failed to create gemini client: %w", err)
	}

	return &GeminiAdapter{
		client: client,
		model:  modelName,
	}, nil
}

func (a *GeminiAdapter) Generate(ctx context.Context, messages []Message) (string, error) {
	model := a.client.GenerativeModel(a.model)
	cs := model.StartChat()

	// Convert history (excluding the last message which is the current prompt)
	var history []*genai.Content
	for i := 0; i < len(messages)-1; i++ {
		history = append(history, &genai.Content{
			Role:  mapRoleToGemini(messages[i].Role),
			Parts: []genai.Part{genai.Text(messages[i].Content)},
		})
	}
	cs.History = history

	// Send the last message
	lastMsg := messages[len(messages)-1].Content
	resp, err := cs.SendMessage(ctx, genai.Text(lastMsg))
	if err != nil {
		return "", fmt.Errorf("gemini generation failed: %w", err)
	}

	if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 {
		return "", fmt.Errorf("empty response from gemini")
	}

	// Extract text from the first part of the first candidate
	if part, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
		return string(part), nil
	}

	return "", fmt.Errorf("unexpected response format")
}

// GenerateStructured reflects the 'target' to create a schema, then unmarshals the result.
func (a *GeminiAdapter) GenerateStructured(ctx context.Context, messages []Message, target any) error {
	t := reflect.TypeOf(target)

	model := a.client.GenerativeModel(a.model)

	// 1. Automatically generate the JSON Schema from the Go struct
	// 1. Recursively map the Go struct to Gemini's Schema format
	model.ResponseMIMEType = "application/json"
	model.ResponseSchema = schemaFromType(t.Elem()) // 2. Convert to the internal genai.Schema format

	var prompt []genai.Part
	for _, m := range messages {
		prompt = append(prompt, genai.Text(m.Content))
	}

	resp, err := model.GenerateContent(ctx, prompt...)
	if err != nil {
		return err
	}

	part := resp.Candidates[0].Content.Parts[0].(genai.Text)
	err = json.Unmarshal([]byte(part), target)
	return err
}

// schemaFromType recursively builds a *genai.Schema tree from a reflect.Type
func schemaFromType(t reflect.Type) *genai.Schema {
	// Follow pointers to the base type
	for t.Kind() == reflect.Ptr {
		t = t.Elem()
	}

	switch t.Kind() {
	case reflect.Struct:
		props := make(map[string]*genai.Schema)
		var required []string

		for i := 0; i < t.NumField(); i++ {
			field := t.Field(i)
			if !field.IsExported() {
				continue
			}

			// Get JSON tag or use field name
			name := field.Tag.Get("json")
			if name == "" || name == "-" {
				name = field.Name
			}

			props[name] = schemaFromType(field.Type)
			required = append(required, name)
		}

		return &genai.Schema{
			Type:       genai.TypeObject,
			Properties: props,
			Required:   required,
		}

	case reflect.Slice, reflect.Array:
		return &genai.Schema{
			Type:  genai.TypeArray,
			Items: schemaFromType(t.Elem()),
		}

	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return &genai.Schema{Type: genai.TypeInteger}

	case reflect.Float32, reflect.Float64:
		return &genai.Schema{Type: genai.TypeNumber}

	case reflect.Bool:
		return &genai.Schema{Type: genai.TypeBoolean}

	default:
		return &genai.Schema{Type: genai.TypeString}
	}
}

func (a *GeminiAdapter) GetProviderName() string {
	return "Google Gemini (" + a.model + ")"
}

// Helper to map your roles to Gemini's expected roles
func mapRoleToGemini(role Role) string {
	switch role {
	case RoleUser:
		return "user"
	case RoleAssistant:
		return "model"
	default:
		return "user" // Gemini doesn't have a specific 'system' role in chat history
	}
}
