package acp

import (
	"bufio"
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"path/filepath"
	"slices"
	"sync"
	"time"

	"albedo/cli/internal/app"
	"albedo/cli/internal/daemon"
	"albedo/cli/internal/daemon/protocol"
)

type Backend interface {
	Models(context.Context) (app.ACPModels, error)
	CurrentModel(context.Context, string) (string, error)
	SetModel(context.Context, string, string) (app.ACPModels, error)
	List(context.Context, *string, *string) (protocol.SessionPage, error)
	History(context.Context, string, daemon.ACPHistoryCursor, int64) ([]daemon.StreamEvent, daemon.ACPHistoryCursor, error)
	Cancel(context.Context, string, string) (protocol.InputCancellation, error)
	New(context.Context, string, []protocol.ACPMCPServer) (protocol.ACPAttachment, error)
	Resume(context.Context, string, string, []protocol.ACPMCPServer) (protocol.ACPAttachment, error)
	Prepare(string, []json.RawMessage) (*daemon.OperationHandle, error)
	Submit(context.Context, *daemon.OperationHandle) error
	Renew(context.Context, string) error
	Release(context.Context, string) error
	Observer(string, protocol.Cursor) *daemon.ChatClient
}
type Connector func(context.Context, string) (Backend, error)

func AppConnector(service *app.Service) Connector {
	return func(ctx context.Context, token string) (Backend, error) { return service.ConnectACP(ctx, token) }
}

type adapter struct {
	backend         Backend
	diagnostics     io.Writer
	ctx             context.Context
	cancel          context.CancelFunc
	writer          *writer
	connect         Connector
	sessions        map[string]*session
	attachmentTails map[string]chan struct{}
	token           string
	workers         sync.WaitGroup
	setupMu         sync.Mutex
	mu              sync.Mutex
}

func Run(ctx context.Context, input io.Reader, output, diagnostics io.Writer, version string, connect Connector) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	token := make([]byte, 32)
	if _, err := rand.Read(token); err != nil {
		return err
	}
	adapter := &adapter{ctx: ctx, cancel: cancel, writer: newWriter(output), connect: connect, token: base64.RawURLEncoding.EncodeToString(token), sessions: map[string]*session{}, diagnostics: diagnostics}
	go adapter.writer.run(ctx)
	lines := make(chan []byte)
	readDone := make(chan error, 1)
	go func() {
		scanner := bufio.NewScanner(input)
		scanner.Buffer(make([]byte, 4096), 100*1024*1024)
		for scanner.Scan() {
			line := slices.Clone(scanner.Bytes())
			select {
			case lines <- line:
			case <-ctx.Done():
				return
			}
		}
		readDone <- scanner.Err()
	}()
	defer adapter.release()
	for {
		select {
		case <-ctx.Done():
			if failure := adapter.writer.failure(); failure != nil {
				return failure
			}
			return ctx.Err()
		case <-adapter.writer.done:
			if err := adapter.writer.failure(); err != nil {
				return err
			}
			return ctx.Err()
		case err := <-readDone:
			// A barrier drains all already queued responses before normal EOF.
			if err != nil {
				return err
			}
			barrier := adapter.writer.flush()
			select {
			case <-barrier:
			case <-adapter.writer.done:
				return adapter.writer.failure()
			case <-ctx.Done():
				return ctx.Err()
			}
			// EOF is a disconnect: accepted work remains daemon-owned.
			return nil
		case line := <-lines:
			message, failure := parse(line)
			if failure != nil {
				if failure.Code == -32602 && message.ID == nil {
					continue
				}
				adapter.respond(message.ID, nil, failure)
				continue
			}
			if message.Method == "" {
				continue
			}
			if message.Method == "initialize" {
				if message.ID != nil {
					adapter.respond(message.ID, initialize(version), nil)
				}
				continue
			}
			if message.Method == "session/cancel" && message.ID == nil {
				var params struct {
					SessionID string `json:"sessionId"`
				}
				if decodeParams(message.Params, &params) == nil {
					adapter.mu.Lock()
					current := adapter.sessions[params.SessionID]
					adapter.mu.Unlock()
					if current != nil {
						current.cancelPrompt()
					}
				}
				continue
			}
			if message.ID == nil {
				continue
			}
			switch message.Method {
			case "authenticate":
				adapter.respond(message.ID, nil, problem(-32602, "invalid_params", "No authentication methods are available"))
			case "session/list":
				adapter.workers.Go(func() { adapter.list(message) })
			case "session/new", "session/resume", "session/load", "session/set_config_option":
				adapter.workers.Go(adapter.setupRequest(message))
			case "session/prompt":
				adapter.prompt(message)
			default:
				adapter.respond(message.ID, nil, problem(-32601, "method_not_found", "Unknown method"))
			}
		}
	}
}
func (a *adapter) response(id json.RawMessage, result any, failure error) map[string]any {
	if id == nil {
		id = json.RawMessage("null")
	}
	message := map[string]any{"jsonrpc": "2.0", "id": id}
	if failure == nil {
		message["result"] = result
		return message
	}
	var rpc *rpcError
	if !errors.As(failure, &rpc) {
		rpc = problem(-32603, "internal_error", failure.Error())
		if api, ok := errors.AsType[*daemon.APIError](failure); ok {
			rpc = daemonProblem(api)
		}
		if _, ok := errors.AsType[*daemon.CompatibilityError](failure); ok {
			rpc = problem(-32603, "daemon_incompatible", "Run `albedo daemon --restart` in a terminal to enable ACP.")
		}
	}
	message["error"] = rpc
	return message
}

// daemonProblem maps a daemon refusal to its ACP error. A missing or unusable
// provider is auth_required; the fix happens in a terminal.
func daemonProblem(api *daemon.APIError) *rpcError {
	switch api.Code {
	case "provider_missing", "provider_unconfigured":
		rpc := problem(-32000, "provider_missing", "No model provider is configured.")
		rpc.Data["instruction"] = "Run `albedo login` in a terminal, then try again."
		return rpc
	case "session_not_found", "not_found":
		return problem(-32002, "session_not_found", api.Message)
	case "daemon_incompatible":
		return problem(-32603, api.Code, api.Message)
	}
	if api.StatusCode >= 500 {
		return problem(-32603, api.Code, api.Message)
	}
	return problem(-32602, api.Code, api.Message)
}
func (a *adapter) respond(id json.RawMessage, result any, failure error) {
	if _, err := a.writer.enqueue(a.response(id, result, failure), false); err != nil {
		a.cancel()
	}
}
func (a *adapter) setup(message request) {
	var params struct {
		Cwd        string            `json:"cwd"`
		SessionID  string            `json:"sessionId"`
		MCPServers []json.RawMessage `json:"mcpServers"`
	}
	if decodeParams(message.Params, &params) != nil || !filepath.IsAbs(params.Cwd) || params.MCPServers == nil && message.Method != "session/resume" {
		a.respond(message.ID, nil, problem(-32602, "invalid_params", "cwd must be absolute and mcpServers must be an array"))
		return
	}
	servers := make([]protocol.ACPMCPServer, 0, len(params.MCPServers))
	for _, raw := range params.MCPServers {
		var fields map[string]json.RawMessage
		var server protocol.ACPMCPServer
		if json.Unmarshal(raw, &fields) != nil || fields["type"] != nil || json.Unmarshal(raw, &server) != nil || server.Name == "" || !filepath.IsAbs(server.Command) || fields["args"] == nil || fields["env"] == nil {
			a.respond(message.ID, nil, problem(-32602, "invalid_params", "Only valid stdio MCP servers are supported"))
			return
		}
		if server.Args == nil {
			server.Args = []string{}
		}
		if server.Env == nil {
			server.Env = []protocol.ACPEnvVariable{}
		}
		servers = append(servers, server)
	}
	if message.Method != "session/new" && !fullID(params.SessionID) {
		a.respond(message.ID, nil, problem(-32602, "invalid_params", "A full session ID is required"))
		return
	}
	if err := a.connectBackend(); err != nil {
		a.respond(message.ID, nil, err)
		return
	}
	// Provider discovery can be slow; do it before acquiring or pausing a lease.
	models, err := a.backend.Models(a.ctx)
	if err != nil {
		a.respond(message.ID, nil, err)
		return
	}
	var attachment protocol.ACPAttachment
	var prior *session
	if message.Method == "session/new" {
		attachment, err = a.backend.New(a.ctx, params.Cwd, servers)
	} else {
		prior = a.pausePrior(params.SessionID)
		attachment, err = a.backend.Resume(a.ctx, params.SessionID, params.Cwd, servers)
		if err != nil {
			if prior != nil {
				prior.resumeObservation()
			}
		}
	}
	if err != nil {
		a.respond(message.ID, nil, err)
		return
	}
	currentModel, err := a.backend.CurrentModel(a.ctx, attachment.SessionID)
	if err != nil {
		if prior != nil {
			prior.resumeObservation()
		} else {
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			defer cancel()
			_ = a.backend.Release(ctx, attachment.SessionID)
		}
		a.respond(message.ID, nil, err)
		return
	}
	models = models.WithCurrent(currentModel)
	a.stopPrior(attachment.SessionID)
	a.mu.Lock()
	current := newSession(a, attachment)
	current.models = models
	if message.Method == "session/load" {
		current.pauseObservation()
		for id := range current.boundaries {
			current.withheld[id] = true
		}
		current.start()
		a.sessions[attachment.SessionID] = current
		a.mu.Unlock()
		if err := current.load(message); err != nil {
			current.stop()
			a.mu.Lock()
			if a.sessions[current.id] == current {
				delete(a.sessions, current.id)
			}
			a.mu.Unlock()
			ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
			defer cancel()
			_ = a.backend.Release(ctx, current.id)
			a.respond(message.ID, nil, err)
		}
		return
	}
	result := map[string]any{"configOptions": modelOptions(models)}
	if message.Method == "session/new" {
		result["sessionId"] = attachment.SessionID
	}
	a.respond(message.ID, result, nil)
	current.start()
	a.sessions[attachment.SessionID] = current
	a.mu.Unlock()
}
func fullID(id string) bool {
	if len(id) != 36 {
		return false
	}
	for i, c := range id {
		if i == 8 || i == 13 || i == 18 || i == 23 {
			if c != '-' {
				return false
			}
		} else if (c < '0' || c > '9') && (c < 'a' || c > 'f') {
			return false
		}
	}
	return true
}
func (a *adapter) prompt(message request) {
	var params struct {
		SessionID string            `json:"sessionId"`
		Prompt    []json.RawMessage `json:"prompt"`
	}
	if err := decodeParams(message.Params, &params); err != nil {
		a.respond(message.ID, nil, problem(-32602, "invalid_params", err.Error()))
		return
	}
	blocks, err := translatePrompt(params.Prompt)
	if err != nil {
		a.respond(message.ID, nil, problem(-32602, "invalid_params", err.Error()))
		return
	}
	a.mu.Lock()
	current := a.sessions[params.SessionID]
	a.mu.Unlock()
	if current == nil {
		a.respond(message.ID, nil, problem(-32002, "session_not_found", "Session is not attached"))
		return
	}
	current.prompt(message, blocks)
}
func (a *adapter) release() {
	a.cancel()
	a.workers.Wait()
	a.mu.Lock()
	defer a.mu.Unlock()
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	for id, current := range a.sessions {
		current.stop()
		_ = a.backend.Release(ctx, id)
	}
}

// Reserve before launching a worker so attachment replacements retain request
// order even when goroutines start out of order. Other sessions remain free.
func (a *adapter) setupRequest(message request) func() {
	if message.Method == "session/new" {
		return func() { a.setup(message) }
	}
	var params struct {
		SessionID string `json:"sessionId"`
	}
	_ = json.Unmarshal(message.Params, &params)
	a.mu.Lock()
	if a.attachmentTails == nil {
		a.attachmentTails = map[string]chan struct{}{}
	}
	prior := a.attachmentTails[params.SessionID]
	done := make(chan struct{})
	a.attachmentTails[params.SessionID] = done
	a.mu.Unlock()
	return func() {
		if prior != nil {
			<-prior
		}
		defer func() {
			a.mu.Lock()
			if a.attachmentTails[params.SessionID] == done {
				delete(a.attachmentTails, params.SessionID)
			}
			close(done)
			a.mu.Unlock()
		}()
		if message.Method == "session/set_config_option" {
			a.setConfigOption(message)
		} else {
			a.setup(message)
		}
	}
}

func (a *adapter) connectBackend() error {
	a.setupMu.Lock()
	defer a.setupMu.Unlock()
	if a.backend != nil {
		return nil
	}
	backend, err := a.connect(a.ctx, a.token)
	if err != nil {
		return err
	}
	a.backend = backend
	return nil
}

// Replace local observation only after the daemon accepts the same-token attachment.
func (a *adapter) stopPrior(id string) {
	a.mu.Lock()
	prior := a.sessions[id]
	if prior != nil {
		prior.mu.Lock()
		prior.lose("observation_uncertain", "Attachment resumed by this connection")
		prior.stop()
		prior.mu.Unlock()
		delete(a.sessions, id)
	}
	a.mu.Unlock()
}

// Reuse the bounded replay buffer during validation. Events already captured by
// a successful replacement must not also be forwarded by the old observer.
func (a *adapter) pausePrior(id string) *session {
	a.mu.Lock()
	prior := a.sessions[id]
	a.mu.Unlock()
	if prior != nil {
		prior.pauseObservation()
	}
	return prior
}