// Agent is a full-featured ACP agent example.
// Run: go run ./examples/agent/main.go
// Then connect a client (see examples/client/main.go) to interact with it.
package main

import (
	"context"
	"fmt"
	"log"
	"sync/atomic"

	"github.com/keepmind9/acp-sdk-go/agent"
	"github.com/keepmind9/acp-sdk-go/core"
	"github.com/keepmind9/acp-sdk-go/helpers"
	"github.com/keepmind9/acp-sdk-go/schema"
)

func ptrBool(b bool) *bool                                 { return &b }
func ptrString(s string) *string                           { return &s }
func ptrSessionId(s string) *schema.SessionId              { v := schema.SessionId(s); return &v }
func ptrStopReason(s schema.StopReason) *schema.StopReason { return &s }

type ExampleAgent struct {
	*agent.Base
	conn         *agent.AgentSideConnection
	sessionCount atomic.Int64
	sessions     map[string]bool
}

func main() {
	impl := &ExampleAgent{
		Base:     &agent.Base{},
		sessions: make(map[string]bool),
	}
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	core.RunAgent(impl)
}

func (a *ExampleAgent) OnConnect(conn *agent.AgentSideConnection) {
	a.conn = conn
	log.Println("client connected")
}

func (a *ExampleAgent) Initialize(_ context.Context, req *schema.InitializeRequest) (*schema.InitializeResponse, error) {
	clientName := ""
	if req.ClientInfo != nil {
		clientName = req.ClientInfo.Name
	}
	log.Printf("initialize: client=%s", clientName)

	protoVers := 1
	return &schema.InitializeResponse{
		ProtocolVersion:   &protoVers,
		AgentCapabilities: a.agentCapabilities(),
		AgentInfo: &schema.Implementation{
			Name:    "example-agent",
			Title:   ptrString("Example ACP Agent"),
			Version: "0.1.0",
		},
	}, nil
}

func (a *ExampleAgent) agentCapabilities() *schema.AgentCapabilities {
	return &schema.AgentCapabilities{
		LoadSession: true,
		PromptCapabilities: &schema.PromptCapabilities{
			Image:           true,
			Audio:           true,
			EmbeddedContext: true,
		},
		SessionCapabilities: &schema.SessionCapabilities{
			List: &schema.SessionListCapabilities{},
		},
	}
}

func (a *ExampleAgent) Authenticate(_ context.Context, req *schema.AuthenticateRequest) (*schema.AuthenticateResponse, error) {
	log.Printf("authenticate: method=%s", req.MethodId)
	return &schema.AuthenticateResponse{}, nil
}

func (a *ExampleAgent) NewSession(_ context.Context, req *schema.NewSessionRequest) (*schema.NewSessionResponse, error) {
	id := a.sessionCount.Add(1)
	sessionID := fmt.Sprintf("sess-%d", id)
	a.sessions[sessionID] = true
	log.Printf("new_session: session=%s cwd=%s", sessionID, req.Cwd)

	update := helpers.AgentMessageUpdate("Hello! I am the example ACP agent.")
	a.conn.SessionUpdate(sessionID, &update)

	return &schema.NewSessionResponse{
		SessionId: ptrSessionId(sessionID),
		Modes:     nil,
	}, nil
}

func (a *ExampleAgent) LoadSession(_ context.Context, req *schema.LoadSessionRequest) (*schema.LoadSessionResponse, error) {
	log.Printf("load_session: session=%s cwd=%s", *req.SessionId, req.Cwd)
	a.sessions[string(*req.SessionId)] = true
	return &schema.LoadSessionResponse{}, nil
}

func (a *ExampleAgent) ListSessions(_ context.Context, req *schema.ListSessionsRequest) (*schema.ListSessionsResponse, error) {
	return &schema.ListSessionsResponse{}, nil
}

func (a *ExampleAgent) Prompt(_ context.Context, req *schema.PromptRequest) (*schema.PromptResponse, error) {
	if !a.sessions[string(*req.SessionId)] {
		a.sessions[string(*req.SessionId)] = true
	}

	log.Printf("prompt: session=%s blocks=%d", *req.SessionId, len(req.Prompt))

	for _, block := range req.Prompt {
		var text string
		switch {
		case block.Text != nil:
			text = block.Text.Text
		case block.Image != nil:
			text = "[image]"
		case block.Audio != nil:
			text = "[audio]"
		case block.Resource_link != nil:
			text = "[resource: " + block.Resource_link.Uri + "]"
		case block.Resource != nil:
			text = "[embedded resource]"
		}

		if text != "" {
			log.Printf("  user: %s", text)
			update := helpers.AgentMessageUpdate(text)
			a.conn.SessionUpdate(*req.SessionId, &update)
		}
	}

	return &schema.PromptResponse{
		StopReason: ptrStopReason(schema.StopReasonEndTurn),
	}, nil
}

func (a *ExampleAgent) Cancel(_ context.Context, notif *schema.CancelNotification) error {
	log.Printf("cancel: session=%s", *notif.SessionId)
	return nil
}

func (a *ExampleAgent) SetSessionMode(_ context.Context, req *schema.SetSessionModeRequest) (*schema.SetSessionModeResponse, error) {
	log.Printf("set_session_mode: session=%s mode=%s", *req.SessionId, *req.ModeId)
	return &schema.SetSessionModeResponse{}, nil
}

func (a *ExampleAgent) SetSessionConfigOption(_ context.Context, req *schema.SetSessionConfigOptionRequest) (*schema.SetSessionConfigOptionResponse, error) {
	log.Printf("set_config_option: session=%s config=%s", *req.SessionId, *req.ConfigId)
	return &schema.SetSessionConfigOptionResponse{}, nil
}

func (a *ExampleAgent) ExtMethod(_ context.Context, name string, params map[string]any) (any, error) {
	log.Printf("ext_method: %s params=%v", name, params)
	return map[string]any{"result": "ok", "method": name}, nil
}

func (a *ExampleAgent) ExtNotification(_ context.Context, name string, params map[string]any) error {
	log.Printf("ext_notification: %s params=%v", name, params)
	return nil
}