package main

import (
    "context"
    "fmt"
    "time"

    llmrouter "github.com/bluefunda/llmrouter"
    "github.com/bluefunda/llmrouter/middleware"
    "github.com/bluefunda/llmrouter/providers/anthropic"
    "github.com/bluefunda/llmrouter/providers/openai"
)

func main() {
    router := llmrouter.New(
        llmrouter.WithProvider("openai", openai.NewFromEnv("openai", "OPENAI_API_KEY")),
        llmrouter.WithProvider("anthropic", anthropic.NewFromEnv()),
        llmrouter.WithMiddleware(
            middleware.Retry(3, time.Second),
            middleware.Timeout(60*time.Second),
        ),
    )

    resp, err := router.Complete(context.Background(), &llmrouter.Request{
        Model: "gpt-4o-mini",
        Messages: []llmrouter.Message{
            {Role: llmrouter.RoleUser, Content: "Hello!"},
        },
    })
    if err != nil {
        panic(err)
    }

    fmt.Println(resp.Choices[0].Message.Content)
}

openai.NewFromEnv("deepseek", "DEEPSEEK_API_KEY")
openai.NewFromEnv("groq", "GROQ_API_KEY")
openai.NewFromEnv("ollama", "")  // no key needed for local
openai.NewFromEnv("sarvam", "SARVAM_API_KEY")

geminiProvider, err := gemini.NewFromEnv()
if err != nil {
    log.Fatal(err)
}
defer router.Close()

router := llmrouter.New(
    llmrouter.WithProvider("openai", openaiProvider),
    llmrouter.WithProvider("anthropic", anthropicProvider),
    llmrouter.WithModelMapping("gpt-4o", "openai"),
    llmrouter.WithModelMapping("claude-sonnet-4-20250514", "anthropic"),
    llmrouter.WithFallback("openai", "anthropic"),
    llmrouter.WithMiddleware(retryMw, cb.Wrap, timeoutMw),
)

middleware.Retry(3, time.Second)
middleware.Retry(3, time.Second, middleware.WithMaxDelay(10*time.Second))
middleware.Retry(3, time.Second, middleware.WithRetryFunc(myRetryPolicy))

cb := middleware.NewCircuitBreaker(5, 30*time.Second)
router := llmrouter.New(
    llmrouter.WithMiddleware(cb.Wrap),
)
fmt.Println(cb.State()) // CBStateClosed / CBStateOpen / CBStateHalfOpen

middleware.Timeout(60 * time.Second)

func Logging(next llmrouter.Provider) llmrouter.Provider {
    return &loggingProvider{Provider: next}
}

router := llmrouter.New(llmrouter.WithMiddleware(Logging))

stream, err := router.Stream(ctx, &llmrouter.Request{
    Model:    "claude-sonnet-4-20250514",
    Messages: []llmrouter.Message{
        {Role: llmrouter.RoleUser, Content: "Write a haiku about Go."},
    },
})
if err != nil {
    log.Fatal(err)
}
defer stream.Close()

for stream.Next() {
    event := stream.Event()
    switch event.Type {
    case llmrouter.EventContentDelta:
        fmt.Print(event.Content)
    case llmrouter.EventToolCallDelta:
        // handle tool call delta
    case llmrouter.EventDone:
        // event.Response holds the final response with usage stats
    }
}
if err := stream.Err(); err != nil {
    log.Fatal(err)
}

tool := llmrouter.Tool{
    Type: "function",
    Function: llmrouter.Function{
        Name:        "get_weather",
        Description: "Get current weather for a location",
        Parameters:  json.RawMessage(`{
            "type": "object",
            "properties": {
                "location": {"type": "string"}
            },
            "required": ["location"]
        }`),
    },
}

resp, _ := router.Complete(ctx, &llmrouter.Request{
    Model:    "gpt-4o-mini",
    Messages: messages,
    Tools:    []llmrouter.Tool{tool},
})

msg := llmrouter.Message{
    Role: llmrouter.RoleUser,
    ContentParts: []llmrouter.ContentPart{
        {Type: "text", Text: "What's in this image?"},
        {Type: "image_url", ImageURL: &llmrouter.ImageURL{URL: "https://..."}},
    },
}

req := &llmrouter.Request{
    Model: "claude-sonnet-4-20250514",
    Messages: []llmrouter.Message{
        {
            Role:         llmrouter.RoleSystem,
            Content:      longSystemPrompt,
            CacheControl: &llmrouter.CacheControl{Type: "ephemeral"},
        },
        {Role: llmrouter.RoleUser, Content: userQuery},
    },
}