package mcp

import (
	"context"

	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"

	"github.com/yan1h/agent-coding-workflow/internal/codeindex"
)

// registerCodeIndexTools registers the pgvector code-search tools. When the
// CodeIndex service is not configured (no embedding endpoint), the tools are not
// registered at all so they don't appear in the agent's tool list.
func registerCodeIndexTools(srv *mcpsdk.Server, deps ServerDeps) {
	if deps.CodeIndex == nil {
		return
	}
	registerSearchCode(srv, deps)
	registerReindexWorkspace(srv, deps)
}

// ===== search_code =====

type searchCodeIn struct {
	WorkspaceID string `json:"workspace_id" jsonschema:"workspace UUID"`
	Query       string `json:"query" jsonschema:"natural-language or code query"`
	TopK        int    `json:"top_k,omitempty" jsonschema:"max results (capped by server)"`
	PathPrefix  string `json:"path_prefix,omitempty" jsonschema:"restrict to a worktree-relative path prefix"`
	Lang        string `json:"lang,omitempty" jsonschema:"restrict to a language label (e.g. go, typescript)"`
}

type codeSnippet struct {
	File      string  `json:"file"`
	StartLine int     `json:"start_line"`
	EndLine   int     `json:"end_line"`
	Snippet   string  `json:"snippet"`
	Score     float32 `json:"score"`
	CommitSHA string  `json:"commit_sha"`
}

type searchCodeOut struct {
	Results []codeSnippet `json:"results"`
}

func registerSearchCode(srv *mcpsdk.Server, deps ServerDeps) {
	mcpsdk.AddTool(srv,
		&mcpsdk.Tool{Name: "search_code", Description: "Semantic code search over the workspace's indexed files; returns ranked file:line snippets."},
		func(ctx context.Context, _ *mcpsdk.CallToolRequest, in searchCodeIn) (*mcpsdk.CallToolResult, searchCodeOut, error) {
			if err := CheckToolFromCtx(ctx, "search_code"); err != nil {
				return nil, searchCodeOut{}, err
			}
			c, err := deps.Caller.Resolve(ctx)
			if err != nil {
				return nil, searchCodeOut{}, err
			}
			w, _, err := scopedWorkspace(ctx, deps, c, in.WorkspaceID)
			if err != nil {
				return nil, searchCodeOut{}, err
			}
			snippets, err := deps.CodeIndex.Search(ctx, w.ID, codeindex.SearchQuery{
				Text: in.Query, TopK: in.TopK, PathPrefix: in.PathPrefix, Lang: in.Lang,
			})
			if err != nil {
				return nil, searchCodeOut{}, err
			}
			out := searchCodeOut{Results: make([]codeSnippet, 0, len(snippets))}
			for _, s := range snippets {
				out.Results = append(out.Results, codeSnippet{
					File: s.FilePath, StartLine: s.StartLine, EndLine: s.EndLine,
					Snippet: s.Content, Score: s.Score, CommitSHA: s.CommitSHA,
				})
			}
			return nil, out, nil
		})
}

// ===== reindex_workspace =====

type reindexWorkspaceIn struct {
	WorkspaceID string `json:"workspace_id" jsonschema:"workspace UUID"`
}

type reindexWorkspaceOut struct {
	OK    bool   `json:"ok"`
	RunID string `json:"run_id"`
}

func registerReindexWorkspace(srv *mcpsdk.Server, deps ServerDeps) {
	mcpsdk.AddTool(srv,
		&mcpsdk.Tool{Name: "reindex_workspace", Description: "Enqueue a fresh code index build at the workspace's current main HEAD."},
		func(ctx context.Context, _ *mcpsdk.CallToolRequest, in reindexWorkspaceIn) (*mcpsdk.CallToolResult, reindexWorkspaceOut, error) {
			if err := CheckToolFromCtx(ctx, "reindex_workspace"); err != nil {
				return nil, reindexWorkspaceOut{}, err
			}
			c, err := deps.Caller.Resolve(ctx)
			if err != nil {
				return nil, reindexWorkspaceOut{}, err
			}
			w, _, err := scopedWorkspace(ctx, deps, c, in.WorkspaceID)
			if err != nil {
				return nil, reindexWorkspaceOut{}, err
			}
			run, err := deps.CodeIndex.EnqueueBuild(ctx, w.ID, nil, "")
			if err != nil {
				return nil, reindexWorkspaceOut{}, err
			}
			return nil, reindexWorkspaceOut{OK: true, RunID: run.ID.String()}, nil
		})
}