//usr/bin/env go run "$0" "$@"; exit
package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	anthropicAPI  = "https://api.anthropic.com/v1/messages/count_tokens"
	model         = "claude-opus-4-8"
	workerCount   = 5 // Number of parallel API workers
	readmePath    = "README.md"
	skillsHeading = "## Available skills"
	todoParagraph = "TODO: describe me"
	maxBodyLines  = 500
)

var httpClient = &http.Client{
	Timeout: 30 * time.Second,
}

type Frontmatter struct {
	Name        string
	Description string
}

// TokenCount breaks down a skill's token usage. Extra is the sum of every
// .md file in the skill folder other than SKILL.md itself.
type TokenCount struct {
	Name        int
	Description int
	Body        int
	Extra       int
	Total       int
}

type SkillInfo struct {
	Dir         string
	Frontmatter Frontmatter
	BodyLines   int
	Tokens      TokenCount
	Errors      []string
}

type TokenJob struct {
	ID   string
	Text string
}

type TokenResult struct {
	ID    string
	Count int
	Err   error
}

func main() {
	workers := flag.Int("workers", workerCount, "Number of parallel API workers")
	flag.Parse()

	apiKey := os.Getenv("ANTHROPIC_API_KEY")
	if apiKey == "" {
		fmt.Fprintln(os.Stderr, "Error: ANTHROPIC_API_KEY environment variable not set")
		os.Exit(1)
	}

	counter := newTokenCounter(apiKey, *workers)
	defer counter.Close()

	skills, err := analyzeSkills(counter)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}

	sort.Slice(skills, func(i, j int) bool {
		return skills[i].Dir < skills[j].Dir
	})

	stats, err := updateReadme(readmePath, skills)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error updating README: %v\n", err)
		os.Exit(1)
	}

	fmt.Fprintf(os.Stderr, "README updated: %d skills maintained, %d errored, %d present in README but missing on disk\n",
		stats.Maintained, stats.Errored, stats.Missing)
}

type TokenCounter struct {
	apiKey  string
	jobs    chan TokenJob
	results chan TokenResult
	wg      sync.WaitGroup
}

func newTokenCounter(apiKey string, workers int) *TokenCounter {
	tc := &TokenCounter{
		apiKey:  apiKey,
		jobs:    make(chan TokenJob, 100),
		results: make(chan TokenResult, 100),
	}
	tc.wg.Add(workers)
	for i := 0; i < workers; i++ {
		go tc.worker()
	}
	go func() {
		tc.wg.Wait()
		close(tc.results)
	}()
	return tc
}

func (tc *TokenCounter) worker() {
	defer tc.wg.Done()
	for job := range tc.jobs {
		count, err := countTokensAPI(tc.apiKey, job.Text)
		tc.results <- TokenResult{ID: job.ID, Count: count, Err: err}
	}
}

func (tc *TokenCounter) Count(id, text string) {
	tc.jobs <- TokenJob{ID: id, Text: text}
}

func (tc *TokenCounter) GetResult() TokenResult {
	result, ok := <-tc.results
	if !ok {
		return TokenResult{Err: fmt.Errorf("results channel closed")}
	}
	return result
}

func (tc *TokenCounter) TryGetResult() (TokenResult, bool) {
	select {
	case result, ok := <-tc.results:
		if !ok {
			return TokenResult{}, false
		}
		return result, true
	default:
		return TokenResult{}, false
	}
}

func (tc *TokenCounter) Close() {
	close(tc.jobs)
}

func analyzeSkills(counter *TokenCounter) ([]SkillInfo, error) {
	skillsDir := "skills"
	entries, err := os.ReadDir(skillsDir)
	if err != nil {
		return nil, fmt.Errorf("cannot read skills directory: %w", err)
	}

	var skills []SkillInfo
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}

		skillPath := filepath.Join(skillsDir, entry.Name())
		skill, err := analyzeSkill(skillPath, counter)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Warning: error analyzing %s: %v\n", entry.Name(), err)
			continue
		}
		skills = append(skills, skill)
	}

	return skills, nil
}

func analyzeSkill(path string, counter *TokenCounter) (SkillInfo, error) {
	skill := SkillInfo{Dir: filepath.Base(path)}

	// Read SKILL.md
	skillMdPath := filepath.Join(path, "SKILL.md")
	content, err := os.ReadFile(skillMdPath)
	if err != nil {
		skill.Errors = append(skill.Errors, fmt.Sprintf("Cannot read SKILL.md: %v", err))
		return skill, nil
	}

	// Parse frontmatter and body
	fm, body, err := parseFrontmatter(string(content))
	if err != nil {
		skill.Errors = append(skill.Errors, fmt.Sprintf("Cannot parse frontmatter: %v", err))
		return skill, nil
	}
	skill.Frontmatter = fm
	trimmedBody := strings.TrimSpace(body)
	if trimmedBody == "" {
		skill.BodyLines = 0
	} else {
		skill.BodyLines = len(strings.Split(trimmedBody, "\n"))
	}

	// Validate
	skill.Errors = append(skill.Errors, validateSkill(skill)...)

	fmt.Fprintf(os.Stderr, "Analyzing %s...\n", skill.Dir)

	// Collect token-counting jobs. SKILL.md is counted as its frontmatter
	// name+description fields plus the body; every other .md file in the
	// folder is counted wholesale.
	var jobs []TokenJob
	jobs = append(jobs, TokenJob{ID: fmt.Sprintf("%s:name", skill.Dir), Text: fm.Name})
	jobs = append(jobs, TokenJob{ID: fmt.Sprintf("%s:description", skill.Dir), Text: fm.Description})
	jobs = append(jobs, TokenJob{ID: fmt.Sprintf("%s:body", skill.Dir), Text: body})

	err = filepath.WalkDir(path, func(p string, d os.DirEntry, walkErr error) error {
		if walkErr != nil {
			skill.Errors = append(skill.Errors, fmt.Sprintf("Cannot walk %s: %v", p, walkErr))
			return nil
		}
		if d.IsDir() {
			return nil
		}
		if !strings.HasSuffix(p, ".md") {
			return nil
		}
		// SKILL.md is already covered by the name/description/body jobs.
		if p == skillMdPath {
			return nil
		}
		fileContent, err := os.ReadFile(p)
		if err != nil {
			skill.Errors = append(skill.Errors, fmt.Sprintf("Cannot read %s: %v", p, err))
			return nil
		}
		relPath, _ := filepath.Rel(path, p)
		jobs = append(jobs, TokenJob{
			ID:   fmt.Sprintf("%s:file:%s", skill.Dir, relPath),
			Text: string(fileContent),
		})
		return nil
	})
	if err != nil {
		skill.Errors = append(skill.Errors, fmt.Sprintf("Walk error: %v", err))
	}

	// Interleave enqueue and drain to avoid backpressure on the buffered channels.
	processResult := func(result TokenResult) {
		if result.Err != nil {
			skill.Errors = append(skill.Errors, fmt.Sprintf("Token count failed for %s: %v", result.ID, result.Err))
			return
		}
		parts := strings.SplitN(result.ID, ":", 3)
		if len(parts) < 2 {
			return
		}
		switch parts[1] {
		case "name":
			skill.Tokens.Name = result.Count
		case "description":
			skill.Tokens.Description = result.Count
		case "body":
			skill.Tokens.Body = result.Count
		case "file":
			skill.Tokens.Extra += result.Count
		}
	}

	outstanding := 0
	for _, job := range jobs {
		counter.Count(job.ID, job.Text)
		outstanding++
		for {
			if result, ok := counter.TryGetResult(); ok {
				processResult(result)
				outstanding--
			} else {
				break
			}
		}
	}
	for outstanding > 0 {
		result := counter.GetResult()
		processResult(result)
		outstanding--
	}

	skill.Tokens.Total = skill.Tokens.Name + skill.Tokens.Description + skill.Tokens.Body + skill.Tokens.Extra

	return skill, nil
}

func parseFrontmatter(content string) (Frontmatter, string, error) {
	lines := strings.Split(content, "\n")
	if len(lines) < 3 || lines[0] != "---" {
		return Frontmatter{}, "", fmt.Errorf("missing frontmatter")
	}

	var fm Frontmatter
	var endIdx int
	var inDescription bool
	var descriptionLines []string

	for i := 1; i < len(lines); i++ {
		if lines[i] == "---" {
			endIdx = i
			break
		}

		line := lines[i]

		// Parse name
		if strings.HasPrefix(line, "name:") {
			fm.Name = strings.TrimSpace(strings.TrimPrefix(line, "name:"))
			continue
		}

		// Parse description (might be multi-line)
		if strings.HasPrefix(line, "description:") {
			descPart := strings.TrimSpace(strings.TrimPrefix(line, "description:"))
			if descPart != "" {
				descriptionLines = append(descriptionLines, descPart)
			}
			inDescription = true
			continue
		}

		// Continue multi-line description
		if inDescription && strings.HasPrefix(line, "  ") {
			descriptionLines = append(descriptionLines, strings.TrimSpace(line))
			continue
		}

		// End of description
		if inDescription && !strings.HasPrefix(line, "  ") {
			inDescription = false
		}
	}

	fm.Description = strings.Join(descriptionLines, " ")

	if endIdx == 0 {
		return Frontmatter{}, "", fmt.Errorf("unclosed frontmatter")
	}

	body := strings.Join(lines[endIdx+1:], "\n")
	return fm, body, nil
}

func validateSkill(skill SkillInfo) []string {
	var errors []string

	// Validate name
	if len(skill.Frontmatter.Name) < 1 || len(skill.Frontmatter.Name) > 64 {
		errors = append(errors, "name must be 1-64 characters")
	}

	namePattern := regexp.MustCompile(`^[a-z0-9]+(-[a-z0-9]+)*$`)
	if !namePattern.MatchString(skill.Frontmatter.Name) {
		errors = append(errors, "name must be lowercase letters, numbers, hyphens only; no leading/trailing/consecutive hyphens")
	}

	if skill.Frontmatter.Name != skill.Dir {
		errors = append(errors, fmt.Sprintf("name '%s' doesn't match directory '%s'", skill.Frontmatter.Name, skill.Dir))
	}

	// Validate description
	if len(skill.Frontmatter.Description) < 1 {
		errors = append(errors, "description is empty")
	} else if len(skill.Frontmatter.Description) > 1024 {
		errors = append(errors, fmt.Sprintf("description is %d characters (max 1024)", len(skill.Frontmatter.Description)))
	}

	// Check body line count
	if skill.BodyLines > maxBodyLines {
		errors = append(errors, fmt.Sprintf("body has %d lines (recommended: < %d)", skill.BodyLines, maxBodyLines))
	}

	return errors
}

func countTokensAPI(apiKey, text string) (int, error) {
	reqBody := map[string]interface{}{
		"model": model,
		"messages": []map[string]string{
			{
				"role":    "user",
				"content": text,
			},
		},
	}

	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return 0, fmt.Errorf("marshal request: %w", err)
	}

	req, err := http.NewRequest("POST", anthropicAPI, bytes.NewBuffer(jsonData))
	if err != nil {
		return 0, fmt.Errorf("create request: %w", err)
	}
	req.Header.Set("x-api-key", apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("anthropic-version", "2023-06-01")

	resp, err := httpClient.Do(req)
	if err != nil {
		return 0, fmt.Errorf("HTTP request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1024))
		return 0, fmt.Errorf("API status %d: %s", resp.StatusCode, string(body))
	}

	var result InputTokens
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return 0, fmt.Errorf("decode response: %w", err)
	}

	return result.InputTokens, nil
}

type InputTokens struct {
	InputTokens int `json:"input_tokens"`
}

// updateStats summarises what updateReadme did.
type updateStats struct {
	Maintained int // skills on disk whose section was (re)written
	Errored    int // skills with one or more errors
	Missing    int // sections in README with no matching skill on disk
}

// parsedSection is one ### skill-name block from the existing README. The
// script rewrites the heading and bullets of every maintained section but
// preserves the paragraph verbatim. For sections whose skill has disappeared
// from disk (orphans), raw is emitted unchanged.
type parsedSection struct {
	name      string
	paragraph []string // prose between heading and first bullet (preserved)
	raw       []string // full original section text (used for orphan passthrough)
}

// updateReadme rewrites the ## Available Skills region of the README: it
// regenerates each section's heading and bullet list, preserving the
// author-written paragraph; inserts a TODO scaffold for skills present on disk
// but not in the README; and leaves sections whose skill is missing from disk
// untouched (with a stderr warning). Everything outside the region is passed
// through byte-for-byte.
func updateReadme(path string, skills []SkillInfo) (updateStats, error) {
	var stats updateStats

	data, err := os.ReadFile(path)
	if err != nil {
		return stats, fmt.Errorf("cannot read %s: %w", path, err)
	}
	lines := strings.Split(string(data), "\n")

	headingIdx := -1
	for i, l := range lines {
		if strings.TrimSpace(l) == skillsHeading {
			headingIdx = i
			break
		}
	}
	if headingIdx == -1 {
		return stats, fmt.Errorf("could not find %q in %s", skillsHeading, path)
	}

	regionEnd := len(lines)
	for i := headingIdx + 1; i < len(lines); i++ {
		if strings.HasPrefix(strings.TrimSpace(lines[i]), "## ") {
			regionEnd = i
			break
		}
	}

	// Leading = lines between the heading and the first ### section.
	firstSectionIdx := regionEnd
	for i := headingIdx + 1; i < regionEnd; i++ {
		if strings.HasPrefix(strings.TrimSpace(lines[i]), "### ") {
			firstSectionIdx = i
			break
		}
	}
	leading := lines[headingIdx+1 : firstSectionIdx]

	// Parse sections and whatever follows them (trailing).
	var existing []parsedSection
	i := firstSectionIdx
	for i < regionEnd {
		if !strings.HasPrefix(strings.TrimSpace(lines[i]), "### ") {
			break // non-section content ends the section run; rest is trailing
		}
		name := strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(lines[i]), "### "))

		bodyStart := i + 1
		bodyEnd := regionEnd
		for j := bodyStart; j < regionEnd; j++ {
			if strings.HasPrefix(strings.TrimSpace(lines[j]), "### ") {
				bodyEnd = j
				break
			}
		}
		raw := lines[i:bodyEnd]
		body := lines[bodyStart:bodyEnd]

		bulletStart := -1
		for k, bl := range body {
			if strings.HasPrefix(strings.TrimSpace(bl), "- ") {
				bulletStart = k
				break
			}
		}
		var paragraph []string
		if bulletStart == -1 {
			paragraph = body
		} else {
			paragraph = body[:bulletStart]
		}
		paragraph = trimLeadingBlanks(trimTrailingBlanks(paragraph))

		existing = append(existing, parsedSection{name: name, paragraph: paragraph, raw: raw})
		i = bodyEnd
	}
	trailing := lines[i:regionEnd]

	existingByPos := existing
	existingByName := make(map[string]parsedSection, len(existing))
	for _, e := range existing {
		existingByName[e.name] = e
	}

	desired := make(map[string]bool, len(skills))
	for _, s := range skills {
		desired[s.Dir] = true
		if len(s.Errors) > 0 {
			stats.Errored++
		}
	}
	stats.Maintained = len(skills)
	for _, e := range existingByPos {
		if !desired[e.name] {
			stats.Missing++
		}
	}

	// Build the rebuilt region.
	var regionLines []string
	lead := trimTrailingBlanks(leading)
	regionLines = append(regionLines, lead...)

	for _, s := range skills {
		paragraph, ok := existingByName[s.Dir]
		if ok && len(paragraph.paragraph) > 0 {
			regionLines = appendBlankSep(regionLines)
			regionLines = append(regionLines, renderSection(s, paragraph.paragraph)...)
		} else {
			// New skill: scaffold with a TODO placeholder.
			if !ok {
				fmt.Fprintf(os.Stderr, "Warning: inserting new section for %s (no description yet)\n", s.Dir)
			}
			regionLines = appendBlankSep(regionLines)
			regionLines = append(regionLines, renderSection(s, []string{todoParagraph})...)
		}
	}

	// Orphans: sections in the README whose skill is gone from disk. Emit
	// their raw text unchanged and warn.
	for _, e := range existingByPos {
		if desired[e.name] {
			continue
		}
		fmt.Fprintf(os.Stderr, "Warning: section %q is in README but missing from disk; leaving untouched\n", e.name)
		regionLines = appendBlankSep(regionLines)
		regionLines = append(regionLines, e.raw...)
	}

	trail := trimLeadingBlanks(trailing)
	if len(trail) > 0 {
		regionLines = appendBlankSep(regionLines)
		regionLines = append(regionLines, trail...)
	}

	// Ensure a blank line separates the region from the following ## heading
	// (if any), matching markdown convention.
	if len(regionLines) > 0 && regionLines[len(regionLines)-1] != "" && regionEnd < len(lines) {
		regionLines = append(regionLines, "")
	}

	var newLines []string
	newLines = append(newLines, lines[:headingIdx]...)
	newLines = append(newLines, lines[headingIdx]) // "## Available Skills"
	newLines = append(newLines, regionLines...)
	newLines = append(newLines, lines[regionEnd:]...)

	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, []byte(strings.Join(newLines, "\n")), 0o644); err != nil {
		return stats, fmt.Errorf("write temp: %w", err)
	}
	if err := os.Rename(tmp, path); err != nil {
		return stats, fmt.Errorf("rename: %w", err)
	}
	return stats, nil
}

// renderSection produces the deterministic heading + (preserved) paragraph +
// bullets + optional warning block. When the skill has errors, all three
// token bullets show 0 and a blockquote warning is appended.
func renderSection(s SkillInfo, paragraph []string) []string {
	var out []string
	out = append(out, "### "+s.Dir, "")
	out = append(out, paragraph...)
	out = append(out, "")

	nameDesc, bodyTk, totalTk := 0, 0, 0
	if len(s.Errors) == 0 {
		nameDesc = s.Tokens.Name + s.Tokens.Description
		bodyTk = s.Tokens.Body
		totalTk = s.Tokens.Total
	}
	out = append(
		out,
		fmt.Sprintf("- Install with: `aubx skills add https://git.secluded.site/agent-skills --skill %s --yes --global`", s.Frontmatter.Name),
		fmt.Sprintf("- Source: [skills/%s/](skills/%s/)", s.Dir, s.Dir),
		fmt.Sprintf("- Name/desc: %d tok", nameDesc),
		fmt.Sprintf("- SKILL.md: %d tok", bodyTk),
		fmt.Sprintf("- Full skill: %d tok (all `.md` files in the folder)", totalTk),
	)

	if len(s.Errors) > 0 {
		out = append(out, "", "> ⚠️ **Errors:**")
		for _, e := range s.Errors {
			out = append(out, "> - "+e)
		}
	}
	return out
}

func appendBlankSep(lines []string) []string {
	if len(lines) == 0 {
		return lines
	}
	if lines[len(lines)-1] != "" {
		return append(lines, "")
	}
	return lines
}

func trimTrailingBlanks(lines []string) []string {
	end := len(lines)
	for end > 0 && strings.TrimSpace(lines[end-1]) == "" {
		end--
	}
	return lines[:end]
}

func trimLeadingBlanks(lines []string) []string {
	start := 0
	for start < len(lines) && strings.TrimSpace(lines[start]) == "" {
		start++
	}
	return lines[start:]
}