package main

import (
    "context"
    "fmt"
    "log"

    "github.com/gollmkit/gollmkit/internal/auth"
    "github.com/gollmkit/gollmkit/internal/config"
    "github.com/gollmkit/gollmkit/internal/providers"
)

func main() {
    // Initialize GoLLMKit
    cfg, err := config.LoadConfig("gollmkit-config.yaml")
    if err != nil {
        log.Fatal(err)
    }

    keyStore, err := auth.NewKeyStoreFromConfig(cfg)
    if err != nil {
        log.Fatal(err)
    }
    defer keyStore.Close()

    rotator := auth.NewKeyRotator(cfg, keyStore)
    validator := auth.NewKeyValidator()
    provider := providers.NewUnifiedProvider(cfg, rotator, validator)

    // Simple completion
    response, err := provider.Invoke(context.Background(), "Explain Go concurrency", providers.RequestOptions{
        Provider:  providers.OpenAI,
        Model:     "gpt-4",
        MaxTokens: 200,
    })

    if err != nil {
        log.Fatal(err)
    }

    fmt.Printf("Response: %s\n", response.Content)
    fmt.Printf("Cost: $%.4f\n", response.Usage.Cost)
}

// Text completion
response, err := provider.Invoke(ctx, "Write a function to sort an array", providers.RequestOptions{
    Provider:    providers.OpenAI,
    Model:       "gpt-4",
    MaxTokens:   500,
    Temperature: 0.3,
})

// Chat conversation
messages := []providers.Message{
    {Role: "system", Content: "You are a Go expert."},
    {Role: "user", Content: "Explain goroutines vs threads."},
}

response, err := provider.Chat(ctx, messages, providers.RequestOptions{
    Provider: providers.Anthropic,
    Model:    "claude-3-sonnet-20240229",
})

response, err := provider.InvokeWithRetry(ctx, prompt, providers.RequestOptions{
    Provider:       providers.OpenAI,
    Model:          "gpt-4",
    RetryAttempts: 3,
    FallbackChain: []providers.ProviderType{providers.Anthropic, providers.Gemini},
})

stream, err := provider.InvokeStream(ctx, "Explain machine learning", providers.RequestOptions{
    Provider: providers.OpenAI,
    Model:    "gpt-4",
    Stream:   true,
})

for chunk := range stream {
    if chunk.Error != nil {
        log.Printf("Stream error: %v", chunk.Error)
        break
    }
    fmt.Print(chunk.Content)
}

// Export usage data
exporter := analytics.NewUsageExporter(rotator)

// Export to CSV
err := exporter.ExportToCSV(ctx, "usage_report.csv", analytics.ExportOptions{
    DateRange: analytics.LastNDays(30),
    Providers: []string{"openai", "anthropic"},
    GroupBy:   "day",
})

// Export to JSON
data, err := exporter.ExportToJSON(ctx, analytics.ExportOptions{
    DateRange: analytics.LastNDays(7),
    Detailed:  true,
})

// Check key health status
healthChecker := auth.NewHealthChecker(keyStore, 5*time.Minute)
healthStatus, err := healthChecker.GetHealthStatus(ctx, map[string][]string{
    "openai": {"primary", "secondary"},
})

for provider, keys := range healthStatus {
    for keyName, isHealthy := range keys {
        fmt.Printf("%s.%s: %s\n", provider, keyName, 
                  map[bool]string{true: "✓ Healthy", false: "✗ Unhealthy"}[isHealthy])
    }
}

// Get detailed usage statistics
stats, err := rotator.GetKeyStatistics(ctx, "openai")
for keyName, usage := range stats {
    fmt.Printf("Key '%s': %d requests, %d tokens, $%.3f\n",
              keyName, usage.UsageCount, usage.TokensUsed, usage.CostUsed)
}

// Provider-specific options
opts := providers.RequestOptions{
    Provider:    providers.OpenAI,
    Model:       "gpt-4",
    Temperature: 0.7,
    MaxTokens:   2000,
    TopP:        0.9,
    SystemPrompt: "You are an expert Go developer.",
}

type RequestOptions struct {
    Provider      ProviderType
    Model         string
    MaxTokens     int
    Temperature   float64
    TopP          float64
    Stream        bool
    SystemPrompt  string
    RetryAttempts int
    FallbackChain []ProviderType
}

type Response struct {
    Content   string
    Model     string
    Usage     Usage
    Provider  ProviderType
    Metadata  map[string]interface{}
}

type Usage struct {
    PromptTokens     int
    CompletionTokens int
    TotalTokens      int
    Cost            float64
}

type UnifiedProvider interface {
    Invoke(ctx context.Context, prompt string, opts RequestOptions) (*Response, error)
    InvokeWithRetry(ctx context.Context, prompt string, opts RequestOptions) (*Response, error)
    Chat(ctx context.Context, messages []Message, opts RequestOptions) (*Response, error)
    InvokeStream(ctx context.Context, prompt string, opts RequestOptions) (<-chan StreamChunk, error)
}

type KeyRotator interface {
    GetNextKey(ctx context.Context, provider string) (*KeySelection, error)
    RecordUsage(ctx context.Context, provider, keyName string, tokens int, cost float64) error
    GetKeyStatistics(ctx context.Context, provider string) (map[string]*KeyUsage, error)
}

func processLLMRequest(ctx context.Context, provider *providers.UnifiedProvider, prompt string) error {
    response, err := provider.InvokeWithRetry(ctx, prompt, providers.RequestOptions{
        Provider:      providers.OpenAI,
        Model:        "gpt-4",
        MaxTokens:    1000,
        RetryAttempts: 3,
        FallbackChain: []providers.ProviderType{providers.Anthropic, providers.Gemini},
    })

    if err != nil {
        return fmt.Errorf("all providers failed: %w", err)
    }

    fmt.Printf("Response: %s\n", response.Content)
    fmt.Printf("Provider: %s, Cost: $%.4f\n", response.Provider, response.Usage.Cost)
    
    return nil
}

// Monitor and export usage
func monitorCosts(ctx context.Context, rotator *auth.KeyRotator) {
    ticker := time.NewTicker(1 * time.Hour)
    defer ticker.Stop()

    for {
        select {
        case <-ticker.C:
            stats, _ := rotator.GetKeyStatistics(ctx, "openai")
            totalCost := 0.0
            for _, usage := range stats {
                totalCost += usage.CostUsed
            }
            
            if totalCost > 400.0 {
                log.Printf("WARNING: Daily cost limit approaching: $%.2f", totalCost)
            }
            
        case <-ctx.Done():
            return
        }
    }
}