package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"blinker/internal/llm"
	"blinker/internal/policy"
	"blinker/internal/tooldef"
)

// A PolicyDraft has the model write a change to a policy file that the
// user asked for. It is a conversation of its own, holding only the
// request and the policies: the main conversation may contain text from
// files, commands or web pages that tries to steer the model, and the
// policy is what confines it. The user reviews every proposal before it
// is saved.
type PolicyDraft struct {
	client   *llm.Client
	loaded   *policy.Loaded
	messages []llm.Message
	callID   string // the proposal waiting for the user's verdict
}

// PolicyProposal is a whole policy file, for the policy Source is now.
type PolicyProposal struct {
	Source policy.Source
	File   *policy.File
}

// maxPolicyAttempts bounds the invalid proposals in a row before giving up.
const maxPolicyAttempts = 3

// policyNames are the names the model proposes a policy file by.
var policyNames = []string{"safe", "unsafe", "agent-safe", "agent-unsafe"}

var proposeTool = llm.Tool{
	Name:        "Propose_Policy",
	Description: "Propose the complete new contents of one policy file. The user sees the change as a diff and decides.",
	Parameters: map[string]any{
		"type": "object",
		"properties": map[string]any{
			"file":   map[string]any{"type": "string", "enum": policyNames, "description": "Which policy file to replace."},
			"policy": map[string]any{"type": "string", "description": "The whole new policy file, comments included."},
		},
		"required": []string{"file", "policy"},
	},
}

// NewPolicyDraft starts a draft of the change the user asked for.
func NewPolicyDraft(client *llm.Client, l *policy.Loaded, env policy.Env, tools []tooldef.Def, request string) *PolicyDraft {
	var entries []string
	for _, t := range tools {
		entries = append(entries, fmt.Sprintf("- %s: %s", t.Name, t.Description))
	}
	toolList := "(none)"
	if len(entries) > 0 {
		toolList = strings.Join(entries, "\n")
	}
	sys := fmt.Sprintf(`You edit the access policies of Blinker, a sandboxed coding agent, as the user asks. Make the smallest change that does what the user asked, and nothing more. Call Propose_Policy with the complete new contents of the one file to change; the user reviews it as a diff. If the request is unclear or needs more than one file changed, say so in words instead of proposing.

The project directory is %s, and the home directory %s.

There are two policies for the project. The safe one ("safe") is what Safe_Bash, and the file tools without asking, get. The unsafe one ("unsafe") is what Unsafe_Bash, and the file tools once the user approves, get. Each starts by including the agent's policy for its mode for every project ("agent-safe" or "agent-unsafe"), and its own rules after that win where they match. Change the project's policies unless the user asks for a change for every project. To let the model do something without asking, change the safe policy; to let it do something once the user approves, change the unsafe one.

A policy file has one rule per line, in three columns separated by spaces: type, access and pattern. The pattern is the rest of the line. Lines starting with # are comments, and the comment lines directly above a rule are its note; keep the comments, and give a rule you add a short note when its purpose is not obvious. For any path, host, tool or variable the last matching rule wins, so a later broad rule undoes earlier narrow ones.
- "include NAME" stands for the rules of the named policy NAME, where the line is.
- "file ACCESS PATH": hide, read or write (write includes read). Paths are globs, relative to the project unless they start with / or ~/, and "." is the project itself: * and ? match within one path component, [abc] one character of a set, and a ** component any number of directories, so go.mod is only the one at the top and **/go.mod is every one. A rule matching a directory covers everything in it. A path that is exactly $NAME, such as $PATH, covers the directories in that environment variable (":"-separated), except home, the project, and directories containing them.
- "host ACCESS HOST": the host is a name, *.example.com (its subdomains only), or *. The access is allow, deny, or ask, which asks the user in the unsafe policy and denies in the safe one. With no matching rule, the host is denied.
- "tool ACCESS NAME": allow or deny. A tool the safe policy allows runs without asking; one only the unsafe policy allows runs once the user approves each call; others never run. The tools are:
%s
- "env ACCESS NAME": whether commands get the environment variables the name matches (a glob: * ? [abc]). The access is allow or deny. With no matching rule only a few basics are passed (PATH, HOME, USER, LANG, LC_*, TERM and the like). API keys, agent sockets and proxy settings are never passed. "env set NAME VALUE" gives a variable a value.
- The project's policies have the line "project %s", which must stay. The agent's have none.

Credentials and shell startup files are blocked whatever the policy says.

The policies now, each with its name for Propose_Policy:
%s`, env.Project, env.Home, toolList, env.Project, policyText(l))
	return &PolicyDraft{
		client: client,
		loaded: l,
		messages: []llm.Message{
			{Role: "system", Content: sys},
			{Role: "user", Content: request},
		},
	}
}

// policyText shows the policies the model may change, by the names it
// proposes them by.
func policyText(l *policy.Loaded) string {
	var b strings.Builder
	srcs := l.Sources()
	for _, name := range policyNames {
		src := srcs[name]
		fmt.Fprintf(&b, "%q, %s (%s):\n%s\n", name, src.Name, src.Path, src.File.Marshal())
	}
	return strings.TrimSuffix(b.String(), "\n")
}

// Propose gets the model's next proposal. check judges a proposal against
// the rest of the policy; an invalid one goes back to the model to fix. If
// the model answers in words instead, that answer is returned in place of
// a proposal.
func (d *PolicyDraft) Propose(ctx context.Context, check func(PolicyProposal) error) (PolicyProposal, string, error) {
	var lastErr error
	for range maxPolicyAttempts {
		res, err := d.client.Complete(ctx, d.messages, []llm.Tool{proposeTool}, nil)
		if err != nil {
			return PolicyProposal{}, "", err
		}
		d.messages = append(d.messages, res.Message())
		if len(res.Calls) == 0 {
			return PolicyProposal{}, res.Reply, nil
		}
		// Every call needs a result; only the first is considered.
		for _, c := range res.Calls[1:] {
			d.messages = append(d.messages, llm.Message{Role: "tool", ToolCallID: c.ID, Content: "Ignored: propose one file at a time."})
		}
		call := res.Calls[0]
		var p PolicyProposal
		if call.Function.Name != proposeTool.Name {
			lastErr = fmt.Errorf("there is no tool named %q", call.Function.Name)
		} else if p, lastErr = d.parse(call.Function.Arguments); lastErr == nil {
			lastErr = check(p)
		}
		if lastErr == nil {
			d.callID = call.ID
			return p, "", nil
		}
		d.messages = append(d.messages, llm.Message{Role: "tool", ToolCallID: call.ID, Content: "Not valid: " + lastErr.Error() + ". Propose again."})
	}
	return PolicyProposal{}, "", fmt.Errorf("the model's proposals were not valid: %w", lastErr)
}

// Revise passes the user's request for changes to the last proposal on to
// the model; Propose then gets the revised one.
func (d *PolicyDraft) Revise(note string) {
	d.messages = append(d.messages, llm.Message{Role: "tool", ToolCallID: d.callID, Content: "The user did not accept this and asks for changes: " + note})
	d.callID = ""
}

// parse reads Propose_Policy's arguments strictly: unknown fields, and a
// policy that is not valid on its own, are errors.
func (d *PolicyDraft) parse(args string) (PolicyProposal, error) {
	var a struct {
		File   string `json:"file"`
		Policy string `json:"policy"`
	}
	dec := json.NewDecoder(bytes.NewReader([]byte(args)))
	dec.DisallowUnknownFields()
	if err := dec.Decode(&a); err != nil {
		return PolicyProposal{}, fmt.Errorf("the arguments: %w", err)
	}
	src, ok := d.loaded.Sources()[a.File]
	if !ok {
		return PolicyProposal{}, errors.New(`"file" must be "safe", "unsafe", "agent-safe" or "agent-unsafe"`)
	}
	f, err := policy.Parse([]byte(a.Policy))
	if err != nil {
		return PolicyProposal{}, err
	}
	if err := src.Check(f); err != nil {
		return PolicyProposal{}, err
	}
	return PolicyProposal{Source: src, File: f}, nil
}