package gemini

import (
	"context"
	"embed"
	"fmt"
	"net/http"
	"net/url"

	"pentagi/pkg/config"
	"pentagi/pkg/providers/pconfig"
	"pentagi/pkg/providers/provider"
	"pentagi/pkg/system"
	"pentagi/pkg/templates"

	"github.com/vxcontrol/langchaingo/httputil"
	"github.com/vxcontrol/langchaingo/llms"
	"github.com/vxcontrol/langchaingo/llms/googleai"
	"github.com/vxcontrol/langchaingo/llms/streaming"
)

//go:embed config.yml models.yml
var configFS embed.FS

const GeminiAgentModel = "gemini-3.1-flash-lite"

const GeminiToolCallIDTemplate = "{r:8:x}"

func BuildProviderConfig(configData []byte) (*pconfig.ProviderConfig, error) {
	defaultOptions := []llms.CallOption{
		llms.WithModel(GeminiAgentModel),
		llms.WithTemperature(1.0),
		llms.WithN(1),
		llms.WithMaxTokens(4000),
	}

	providerConfig, err := pconfig.LoadConfigData(configData, defaultOptions)
	if err != nil {
		return nil, err
	}

	return providerConfig, nil
}

func DefaultProviderConfig() (*pconfig.ProviderConfig, error) {
	configData, err := configFS.ReadFile("config.yml")
	if err != nil {
		return nil, err
	}

	return BuildProviderConfig(configData)
}

func DefaultModels() (pconfig.ModelsConfig, error) {
	configData, err := configFS.ReadFile("models.yml")
	if err != nil {
		return nil, err
	}

	return pconfig.LoadModelsConfigData(configData)
}

type geminiProvider struct {
	llm            *googleai.GoogleAI
	models         pconfig.ModelsConfig
	providerName   provider.ProviderName
	providerConfig *pconfig.ProviderConfig
}

func New(
	cfg *config.Config,
	providerName provider.ProviderName,
	providerConfig *pconfig.ProviderConfig,
) (provider.Provider, error) {
	opts := []googleai.Option{
		googleai.WithRest(),
		googleai.WithAPIKey(cfg.GeminiAPIKey),
		googleai.WithDefaultModel(GeminiAgentModel),
	}

	if _, err := url.Parse(cfg.GeminiServerURL); err != nil {
		return nil, fmt.Errorf("failed to parse Gemini server URL: %w", err)
	}

	httpClient, err := newHTTPClient(cfg)
	if err != nil {
		return nil, err
	}

	opts = append(opts, googleai.WithHTTPClient(httpClient))

	models, err := DefaultModels()
	if err != nil {
		return nil, err
	}

	client, err := googleai.New(context.Background(), opts...)
	if err != nil {
		return nil, err
	}

	return &geminiProvider{
		llm:            client,
		models:         models,
		providerName:   providerName,
		providerConfig: providerConfig,
	}, nil
}

func newHTTPClient(cfg *config.Config) (*http.Client, error) {
	httpClient, err := system.GetHTTPClient(cfg)
	if err != nil {
		return nil, fmt.Errorf("failed to create HTTP client: %w", err)
	}

	httpClient.Transport = &httputil.ApiKeyTransport{
		Transport: httpClient.Transport,
		APIKey:    cfg.GeminiAPIKey,
		BaseURL:   cfg.GeminiServerURL,
	}

	return httpClient, nil
}

func (p *geminiProvider) Type() provider.ProviderType {
	return provider.ProviderGemini
}

func (p *geminiProvider) Name() provider.ProviderName {
	return p.providerName
}

func (p *geminiProvider) GetRawConfig() []byte {
	return p.providerConfig.GetRawConfig()
}

func (p *geminiProvider) GetProviderConfig() *pconfig.ProviderConfig {
	return p.providerConfig
}

func (p *geminiProvider) GetPriceInfo(opt pconfig.ProviderOptionsType) *pconfig.PriceInfo {
	return p.providerConfig.GetPriceInfoForType(opt)
}

func (p *geminiProvider) GetModels() pconfig.ModelsConfig {
	return p.models
}

func (p *geminiProvider) Model(opt pconfig.ProviderOptionsType) string {
	model := GeminiAgentModel
	opts := llms.CallOptions{Model: &model}
	for _, option := range p.providerConfig.GetOptionsForType(opt) {
		option(&opts)
	}

	return opts.GetModel()
}

func (p *geminiProvider) ModelWithPrefix(opt pconfig.ProviderOptionsType) string {
	// Gemini provider doesn't need prefix support (passthrough mode in LiteLLM)
	return p.Model(opt)
}

func (p *geminiProvider) Call(
	ctx context.Context,
	opt pconfig.ProviderOptionsType,
	prompt string,
) (string, error) {
	return provider.WrapGenerateFromSinglePrompt(
		ctx, p, opt, p.llm, prompt,
		p.providerConfig.GetOptionsForType(opt)...,
	)
}

func (p *geminiProvider) CallEx(
	ctx context.Context,
	opt pconfig.ProviderOptionsType,
	chain []llms.MessageContent,
	streamCb streaming.Callback,
) (*llms.ContentResponse, error) {
	return provider.WrapGenerateContent(
		ctx, p, opt, p.llm.GenerateContent, chain,
		append([]llms.CallOption{
			llms.WithStreamingFunc(streamCb),
		}, p.providerConfig.GetOptionsForType(opt)...)...,
	)
}

func (p *geminiProvider) CallWithTools(
	ctx context.Context,
	opt pconfig.ProviderOptionsType,
	chain []llms.MessageContent,
	tools []llms.Tool,
	streamCb streaming.Callback,
) (*llms.ContentResponse, error) {
	return provider.WrapGenerateContent(
		ctx, p, opt, p.llm.GenerateContent, chain,
		append([]llms.CallOption{
			llms.WithTools(tools),
			llms.WithStreamingFunc(streamCb),
		}, p.providerConfig.GetOptionsForType(opt)...)...,
	)
}

func (p *geminiProvider) CallWithExtraOptions(
	ctx context.Context,
	opt pconfig.ProviderOptionsType,
	chain []llms.MessageContent,
	tools []llms.Tool,
	streamCb streaming.Callback,
	extra ...llms.CallOption,
) (*llms.ContentResponse, error) {
	options := []llms.CallOption{llms.WithStreamingFunc(streamCb)}
	if len(tools) > 0 {
		options = append(options, llms.WithTools(tools))
	}
	options = append(options, p.providerConfig.GetOptionsForType(opt)...)
	options = append(options, extra...)

	return provider.WrapGenerateContent(ctx, p, opt, p.llm.GenerateContent, chain, options...)
}

func (p *geminiProvider) GetUsage(info map[string]any) pconfig.CallUsage {
	return pconfig.NewCallUsage(info)
}

func (p *geminiProvider) GetToolCallIDTemplate(ctx context.Context, prompter templates.Prompter) (string, error) {
	return provider.DetermineToolCallIDTemplate(ctx, p, pconfig.OptionsTypeSimple, prompter, GeminiToolCallIDTemplate)
}