package main

import (
    "context"
    "log"
    "strings"

    "github.com/CoolBanHub/aggo/agent"
    "github.com/CoolBanHub/aggo/model"
    "github.com/CoolBanHub/aggo/memory"
    "github.com/CoolBanHub/aggo/memory/builtin"
    "github.com/CoolBanHub/aggo/memory/builtin/storage"
    "github.com/cloudwego/eino/adk"
    "github.com/cloudwego/eino/schema"
)

func main() {
    ctx := context.Background()

    // 创建聊天模型
    cm, _ := model.NewChatModel(
        model.WithBaseUrl("https://api.openai.com/v1"),
        model.WithAPIKey("your-api-key"),
        model.WithModel("gpt-4"),
    )

    // 创建记忆 provider
    provider, _ := memory.GlobalRegistry().CreateProvider("builtin", &builtin.ProviderConfig{
        ChatModel: cm,
        Storage:   storage.NewMemoryStore(),
        MemoryConfig: &builtin.MemoryConfig{
            EnableUserMemories:   true,
            EnableSessionSummary: true,
            MemoryLimit:          10,
            Retrieval:            builtin.RetrievalLastN,
        },
    })
    defer provider.Close()

    ag, _ := agent.NewAgentBuilder(cm).
        WithInstruction("你是一个友好的 AI 助手").
        WithMemory(provider).
        Build(ctx)

    runner := adk.NewTypedRunner(adk.TypedRunnerConfig[*schema.AgenticMessage]{Agent: ag})
    iter := runner.Run(ctx, []*schema.AgenticMessage{
        schema.UserAgenticMessage("你好，介绍一下你自己"),
    }, adk.WithSessionValues(map[string]any{
        "userID":    "demo-user",
        "sessionID": "demo-session",
    }))

    for {
        event, ok := iter.Next()
        if !ok {
            break
        }
        if event.Err != nil {
            log.Fatal(event.Err)
        }
        if event.Output != nil && event.Output.MessageOutput != nil {
            if msg, err := event.Output.MessageOutput.GetMessage(); err == nil && msg != nil {
                log.Printf("AI: %s", assistantText(msg))
            }
        }
    }
}

func assistantText(msg *schema.AgenticMessage) string {
    var b strings.Builder
    for _, block := range msg.ContentBlocks {
        if block != nil && block.AssistantGenText != nil {
            b.WriteString(block.AssistantGenText.Text)
        }
    }
    return b.String()
}

ag, err := agent.NewAgentBuilder(chatModel).
    WithInstruction("你是一个AI助手").
    WithMemory(provider).
    WithTools(tools...).
    WithMaxStep(10).
    Build(ctx)

provider, err := memory.GlobalRegistry().CreateProvider("builtin", &builtin.ProviderConfig{
    ChatModel: chatModel,
    Storage:   storage.NewMemoryStore(),
    MemoryConfig: &builtin.MemoryConfig{
        EnableSessionSummary: true,               // 启用会话摘要
        EnableUserMemories:   true,               // 启用用户长期记忆
        MemoryLimit:          10,                 // 历史消息条数限制
        Retrieval:            builtin.RetrievalLastN, // 检索策略
    },
})
if err != nil {
    panic(err)
}

import (
    "github.com/CoolBanHub/aggo/memory"
    "github.com/CoolBanHub/aggo/memory/mem0"
)

provider, err := memory.GlobalRegistry().CreateProvider("mem0", &mem0.ProviderConfig{
    BaseURL:           "https://api.mem0.ai",
    APIKey:            "your-mem0-api-key",
    Mode:              mem0.ModeHosted,
    SearchMsgLimit:    6,
    SearchResultLimit: 5,
    OutputMemoryLimit: 5,
})
if err != nil {
    panic(err)
}

import (
    cronPkg "github.com/CoolBanHub/aggo/cron"
    cronTool "github.com/CoolBanHub/aggo/tools/cron"
)

store := cronPkg.NewFileStore("/path/to/cron_jobs.json")
service := cronPkg.NewCronService(store, nil)
cronTools := cronTool.GetTools(service)

result, _ := cronPkg.NewCronAgent(ctx, chatModel, cronTools,
    cronPkg.WithMaxJobsPerUser(20),
)

_ = result.Service.Start()
defer result.Service.Stop()

result, _ := cronPkg.NewCronAgent(ctx, chatModel, cronTools,
    cronPkg.WithOnJobTriggered(func(job *cronPkg.CronJob) {
        fmt.Printf("任务触发: %s\n", job.Payload.Message)
    }),
)

import (
    "github.com/CoolBanHub/aggo/database/milvus"
    "github.com/milvus-io/milvus/client/v2/milvusclient"
)

// 创建 Milvus 客户端
client, _ := milvusclient.New(ctx, &milvusclient.ClientConfig{
    Address: "127.0.0.1:19530",
    DBName:  "",  // 使用默认数据库
})

// 创建向量数据库实例（内部使用 eino-ext milvus2 组件）
vectorDB, _ := milvus.NewMilvus(ctx, milvus.MilvusConfig{
    Client:         client,
    CollectionName: "knowledge_vectors",
    EmbeddingDim:   1024,
    Embedding:      embeddingModel,
})

import "github.com/CoolBanHub/aggo/database/postgres"

vectorDB, _ := postgres.NewPostgres(postgres.PostgresConfig{
    Client:          gormDB,  // GORM 数据库实例
    CollectionName:  "knowledge_vectors",
    VectorDimension: 1024,
    Embedding:       embeddingModel,
})

import "github.com/CoolBanHub/aggo/model"

chatModel, _ := model.NewChatModel(
    model.WithBaseUrl("https://api.openai.com/v1"),
    model.WithAPIKey("your-api-key"),
    model.WithModel("gpt-4"),
    model.WithReasoningEffortLevel("medium"),  // 推理强度: low, medium, high
)

embeddingModel, _ := model.NewEmbModel(
    model.WithBaseUrl("https://api.openai.com/v1"),
    model.WithAPIKey("your-api-key"),
    model.WithModel("text-embedding-3-large"),
    model.WithDimensions(1024),
)

import "github.com/CoolBanHub/aggo/tools"

knowledgeTools := tools.GetKnowledgeTools(vectorDB, retriever, &retriever.Options{
    TopK:           utils.ValueToPtr(10),
    ScoreThreshold: utils.ValueToPtr(0.1),
})

import (
    "github.com/CoolBanHub/aggo/tools"
    "github.com/CoolBanHub/aggo/tools/database"
)

// 默认只允许 SELECT/SHOW/DESCRIBE/EXPLAIN/PRAGMA/WITH 等只读查询
dbTools := tools.GetDatabaseTools(gormDB)

// 如确需写操作，必须显式开启
writeTools := tools.GetDatabaseTools(gormDB, database.WithAllowWrite(true))

import (
    "github.com/CoolBanHub/aggo/tools"
    "github.com/CoolBanHub/aggo/tools/shell"
)

// 默认限制工作目录并拒绝高危命令
shellTools := tools.GetShellTools()

// 生产环境建议进一步配置命令白名单
restrictedShellTools := tools.GetShellTools(shell.WithAllowedCommands("ls", "pwd", "cat"))

import (
    "net/http"

    "github.com/CoolBanHub/aggo/pkg/adapter"
    "github.com/CoolBanHub/aggo/pkg/sse"
    "github.com/cloudwego/eino/schema"
)

// 创建 SSE 写入器
writer := sse.NewWriter("session-id", w)
if writer == nil {
    http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
    return
}
defer writer.Close()

// 将模型流式输出写为 SSE；Writer.Stream 会在正常结束时写出 [DONE]
stream, err := chatModel.Stream(ctx, []*schema.AgenticMessage{
    schema.UserAgenticMessage("你好"),
})
if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
}

_ = writer.Stream(ctx, stream, func(output *schema.AgenticMessage, index int) any {
    return adapter.MessageToOpenaiStreamResponse(output, index)
})

defer provider.Close()  // 确保在创建 provider 后立即添加 defer