package agent

import (
	"context"
	"io"
	"log/slog"

	"github.com/keepmind9/acp-sdk-go/rpc"
	"github.com/keepmind9/acp-sdk-go/schema"
)

func ptrSessionId(s string) *schema.SessionId { v := schema.SessionId(s); return &v }

// AgentSideConnection wraps an rpc.Connection and provides typed methods
// for the agent to call on the client (session/update, request_permission, etc.).
// It also dispatches incoming JSON-RPC messages to the Agent implementation.
type AgentSideConnection struct {
	conn *rpc.Connection
}

// NewAgentSideConnection creates a new agent-side connection.
// The agent parameter implements the Agent interface to handle incoming calls.
// The connection reads from reader and writes to writer.
func NewAgentSideConnection(agent Agent, reader io.Reader, writer io.Writer) *AgentSideConnection {
	handler := buildAgentRouter(agent)
	conn := rpc.NewConnection(context.Background(), handler, reader, writer)
	asc := &AgentSideConnection{conn: conn}

	// Notify agent of connection if it implements OnConnectHandler
	if h, ok := agent.(OnConnectHandler); ok {
		h.OnConnect(asc)
	}

	return asc
}

// ReceiveLoop starts the main event loop that reads and processes incoming messages.
// This blocks until the connection is closed or an error occurs.
func (c *AgentSideConnection) ReceiveLoop() {
	c.conn.ReceiveLoop()
}

// Close shuts down the connection.
func (c *AgentSideConnection) Close() error {
	return c.conn.Close()
}

// SetLogger configures the logger for this connection.
func (c *AgentSideConnection) SetLogger(logger *slog.Logger) {
	c.conn.SetLogger(logger)
}

// --- Client methods (agent calls these to communicate with the client) ---

// SessionUpdate sends a session update notification to the client.
func (c *AgentSideConnection) SessionUpdate(sessionID string, update *schema.SessionUpdate) error {
	notif := &schema.SessionNotification{
		SessionId: ptrSessionId(sessionID),
	}
	if update != nil {
		notif.Update = update
	}
	return c.conn.SendNotification(schema.MethodSessionUpdate, notif)
}

// RequestPermission sends a permission request to the client and waits for the response.
func (c *AgentSideConnection) RequestPermission(req *schema.RequestPermissionRequest) (*schema.RequestPermissionResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodSessionRequest_permission, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.RequestPermissionResponse](result)
}

// ReadTextFile requests the client to read a text file.
func (c *AgentSideConnection) ReadTextFile(req *schema.ReadTextFileRequest) (*schema.ReadTextFileResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodFsRead_text_file, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.ReadTextFileResponse](result)
}

// WriteTextFile requests the client to write a text file.
func (c *AgentSideConnection) WriteTextFile(req *schema.WriteTextFileRequest) (*schema.WriteTextFileResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodFsWrite_text_file, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.WriteTextFileResponse](result)
}

// CreateTerminal requests the client to create a terminal.
func (c *AgentSideConnection) CreateTerminal(req *schema.CreateTerminalRequest) (*schema.CreateTerminalResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodTerminalCreate, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.CreateTerminalResponse](result)
}

// TerminalOutput requests terminal output from the client.
func (c *AgentSideConnection) TerminalOutput(req *schema.TerminalOutputRequest) (*schema.TerminalOutputResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodTerminalOutput, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.TerminalOutputResponse](result)
}

// ReleaseTerminal requests the client to release a terminal.
func (c *AgentSideConnection) ReleaseTerminal(req *schema.ReleaseTerminalRequest) (*schema.ReleaseTerminalResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodTerminalRelease, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.ReleaseTerminalResponse](result)
}

// WaitForTerminalExit waits for a terminal to exit.
func (c *AgentSideConnection) WaitForTerminalExit(req *schema.WaitForTerminalExitRequest) (*schema.WaitForTerminalExitResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodTerminalWait_for_exit, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.WaitForTerminalExitResponse](result)
}

// KillTerminal requests the client to kill a terminal.
func (c *AgentSideConnection) KillTerminal(req *schema.KillTerminalRequest) (*schema.KillTerminalResponse, error) {
	result, err := c.conn.SendRequest(context.Background(), schema.MethodTerminalKill, req)
	if err != nil {
		return nil, err
	}
	return rpc.DecodeResponse[*schema.KillTerminalResponse](result)
}

// ExtMethod sends an extension request to the client.
func (c *AgentSideConnection) ExtMethod(method string, params map[string]any) (any, error) {
	return c.conn.SendRequest(context.Background(), "_"+method, params)
}

// ExtNotification sends an extension notification to the client.
func (c *AgentSideConnection) ExtNotification(method string, params map[string]any) error {
	return c.conn.SendNotification("_"+method, params)
}