package main

import (
	"context"
	"encoding/json"
	"sort"

	"github.com/mark3labs/mcp-go/mcp"
)

type listDomainsArgs struct{}

type domainAgg struct {
	Domain    string         `json:"domain"`
	Functions int            `json:"functions"`
	Types     int            `json:"types"`
	Pure      int            `json:"pure"`
	Impure    int            `json:"impure"`
	Pipelines int            `json:"pipelines"`
	Component int            `json:"component"`
	ByLang    map[string]int `json:"by_lang"`
}

func listDomainsTool() mcp.Tool {
	return mcp.NewTool("fn_list_domains",
		mcp.WithDescription("List all registry domains with counts: functions, types, purity breakdown, language breakdown. Useful to scope a search."),
	)
}

func (d *deps) handleListDomains(ctx context.Context, _ mcp.CallToolRequest, _ listDomainsArgs) (*mcp.CallToolResult, error) {
	conn := d.db.Conn()
	agg := map[string]*domainAgg{}
	get := func(dom string) *domainAgg {
		a := agg[dom]
		if a == nil {
			a = &domainAgg{Domain: dom, ByLang: map[string]int{}}
			agg[dom] = a
		}
		return a
	}

	rows, err := conn.Query(`SELECT domain, kind, purity, lang FROM functions`)
	if err != nil {
		return mcp.NewToolResultError("query functions: " + err.Error()), nil
	}
	for rows.Next() {
		var dom, kind, purity, lang string
		if err := rows.Scan(&dom, &kind, &purity, &lang); err != nil {
			rows.Close()
			return mcp.NewToolResultError("scan functions: " + err.Error()), nil
		}
		a := get(dom)
		a.Functions++
		a.ByLang[lang]++
		switch purity {
		case "pure":
			a.Pure++
		case "impure":
			a.Impure++
		}
		switch kind {
		case "pipeline":
			a.Pipelines++
		case "component":
			a.Component++
		}
	}
	rows.Close()

	rows, err = conn.Query(`SELECT domain, lang FROM types`)
	if err != nil {
		return mcp.NewToolResultError("query types: " + err.Error()), nil
	}
	for rows.Next() {
		var dom, lang string
		if err := rows.Scan(&dom, &lang); err != nil {
			rows.Close()
			return mcp.NewToolResultError("scan types: " + err.Error()), nil
		}
		a := get(dom)
		a.Types++
		a.ByLang["type:"+lang]++
	}
	rows.Close()

	out := make([]*domainAgg, 0, len(agg))
	for _, a := range agg {
		out = append(out, a)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Domain < out[j].Domain })

	b, _ := json.MarshalIndent(map[string]any{"domains": out}, "", "  ")
	return mcp.NewToolResultText(string(b)), nil
}
