main.go

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"

	"path/filepath"

	"github.com/devthedeveloper/codebase-agent/internal/ollama"
	"github.com/devthedeveloper/codebase-agent/internal/pentest"
	"github.com/devthedeveloper/codebase-agent/internal/scanner"
	"github.com/devthedeveloper/codebase-agent/internal/security"
	"github.com/devthedeveloper/codebase-agent/internal/understand"
)

const version = "0.1.0"
const maxContextChars = 120_000 // ~30k tokens

const systemPrompt = `You are a codebase onboarding expert. You help developers understand unfamiliar codebases quickly and accurately.

You have been given the full context of a repository including file structure and source code. Use this to answer questions about:
- Architecture and design patterns
- How specific features work
- Where to find specific functionality
- Data flow and control flow
- Dependencies and relationships between components
- Where to add new features or make changes

Guidelines:
- Always cite file paths and line numbers when referencing code
- Be concise but thorough
- If you're unsure, say so and suggest what to look at
- Think about the codebase holistically, not just individual files`

func main() {
	if len(os.Args) < 2 {
		printUsage()
		os.Exit(1)
	}

	switch os.Args[1] {
	case "ask":
		if len(os.Args) < 3 {
			fmt.Println("Usage: codebase-agent ask <repo-path> [-q \"question\"]")
			os.Exit(1)
		}
		cmdAsk(os.Args[2], os.Args[3:])

	case "understand":
		if len(os.Args) < 3 {
			fmt.Println("Usage: codebase-agent understand <repo-path> [--exclude dir1,dir2] [--include dir1,dir2]")
			os.Exit(1)
		}
		cmdUnderstand(os.Args[2], os.Args[3:])

	case "pentest":
		if len(os.Args) < 3 {
			fmt.Println("Usage: codebase-agent pentest <repo-path> [--local] [--exclude dir1,dir2]")
			os.Exit(1)
		}
		cmdPentest(os.Args[2], os.Args[3:])

	case "security":
		if len(os.Args) < 3 {
			fmt.Println("Usage: codebase-agent security <repo-path> [--local] [--exclude dir1,dir2]")
			os.Exit(1)
		}
		cmdSecurity(os.Args[2], os.Args[3:])

	case "scan":
		if len(os.Args) < 3 {
			fmt.Println("Usage: codebase-agent scan <repo-path>")
			os.Exit(1)
		}
		cmdScan(os.Args[2])

	case "version", "--version", "-v":
		fmt.Printf("codebase-agent v%s\n", version)

	case "help", "--help", "-h":
		printUsage()

	default:
		fmt.Printf("Unknown command: %s\n\n", os.Args[1])
		printUsage()
		os.Exit(1)
	}
}

func printUsage() {
	fmt.Printf(`codebase-agent v%s — Understand any codebase in minutes

Usage:
  codebase-agent pentest <repo-path>                            AI penetration test with PoC exploits
  codebase-agent security <repo-path>                           Security scan + threat model + vuln report
  codebase-agent understand <repo-path>                        Generate docs + interactive HTML
  codebase-agent understand <repo-path> --exclude dir1,dir2   Skip specific directories
  codebase-agent understand <repo-path> --include src,lib     Only scan specific directories
  codebase-agent ask <repo-path>                              Interactive Q&A
  codebase-agent ask <repo-path> -q "question"                Single question
  codebase-agent scan <repo-path>                             Show repo stats

Flags:
  --local        Use local Ollama instead of cloud (requires ollama + model)
  --exclude      Comma-separated directories to skip
  --include      Comma-separated directories to scan (only these)

Examples:
  codebase-agent pentest ./my-project
  codebase-agent security ./my-project
  codebase-agent understand ./my-project
  codebase-agent ask ./my-project -q "How does auth work?"
`, version)
}

func cmdUnderstand(repoPath string, args []string) {
	// Parse flags
	var opts *scanner.ScanOptions
	useLocal := false
	for i, arg := range args {
		if opts == nil {
			opts = &scanner.ScanOptions{}
		}
		if arg == "--exclude" && i+1 < len(args) {
			opts.ExcludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--include" && i+1 < len(args) {
			opts.IncludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--local" {
			useLocal = true
		}
	}

	fmt.Printf("⏳ Scanning %s...\n", repoPath)
	if opts != nil && len(opts.ExcludeDirs) > 0 {
		fmt.Printf("   Excluding: %s\n", strings.Join(opts.ExcludeDirs, ", "))
	}
	if opts != nil && len(opts.IncludeDirs) > 0 {
		fmt.Printf("   Including only: %s\n", strings.Join(opts.IncludeDirs, ", "))
	}

	summary, err := scanner.ScanRepoWithOptions(repoPath, opts)
	if err != nil {
		fmt.Printf("❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("✅ Loaded %d files (%d lines) across %d languages\n\n",
		summary.TotalFiles, summary.TotalLines, len(summary.Languages))

	absPath, _ := filepath.Abs(repoPath)
	outDir := filepath.Join(absPath, "understanding")

	fmt.Printf("📝 Generating code understanding docs → %s/\n\n", outDir)

	gen := understand.NewGenerator(summary, outDir, useLocal)
	if useLocal {
		fmt.Println("   Using local Ollama")
	} else {
		fmt.Println("   Using Ollama Cloud")
	}
	if err := gen.Generate(); err != nil {
		fmt.Printf("\n❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\n🎉 Done! Check the understanding/ directory:\n")
	fmt.Printf("   📄 understanding/INDEX.md        — Start here\n")
	fmt.Printf("   📄 understanding/README.md       — Project overview\n")
	fmt.Printf("   📄 understanding/architecture.md — System diagrams\n")
	fmt.Printf("   📄 understanding/data-flow.md    — Data flow analysis\n")
	fmt.Printf("   📄 understanding/module-map.md   — Module relationships\n")
	fmt.Printf("   🌐 understanding/index.html      — Interactive website\n")
	fmt.Printf("\n   Open in browser: file://%s/index.html\n", outDir)
}

func cmdPentest(repoPath string, args []string) {
	var opts *scanner.ScanOptions
	useLocal := false
	for i, arg := range args {
		if opts == nil {
			opts = &scanner.ScanOptions{}
		}
		if arg == "--exclude" && i+1 < len(args) {
			opts.ExcludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--include" && i+1 < len(args) {
			opts.IncludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--local" {
			useLocal = true
		}
	}

	fmt.Printf("⚔️  Penetration test: %s\n", repoPath)

	summary, err := scanner.ScanRepoWithOptions(repoPath, opts)
	if err != nil {
		fmt.Printf("❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("✅ Loaded %d files (%d lines) across %d languages\n\n",
		summary.TotalFiles, summary.TotalLines, len(summary.Languages))

	absPath, _ := filepath.Abs(repoPath)
	outDir := filepath.Join(absPath, "pentest-report")

	fmt.Printf("📝 Generating pentest report → %s/\n\n", outDir)

	gen := pentest.NewGenerator(summary, outDir, useLocal)
	if err := gen.Generate(); err != nil {
		fmt.Printf("\n❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\n⚔️  Pentest report complete!\n")
	fmt.Printf("   💉 pentest-report/injection-report.md      — Injection attacks + PoC payloads\n")
	fmt.Printf("   🔓 pentest-report/auth-attack-report.md    — Auth bypass + session attacks\n")
	fmt.Printf("   📡 pentest-report/data-exposure-report.md  — Data leakage + extraction\n")
	fmt.Printf("   🧠 pentest-report/business-logic-report.md — Logic flaws + race conditions\n")
	fmt.Printf("   🌐 pentest-report/index.html               — Interactive report\n")
	fmt.Printf("\n   ⚠️  CONFIDENTIAL — Contains exploit payloads\n")
	fmt.Printf("   Open in browser: file://%s/index.html\n", outDir)
}

func cmdSecurity(repoPath string, args []string) {
	var opts *scanner.ScanOptions
	useLocal := false
	for i, arg := range args {
		if opts == nil {
			opts = &scanner.ScanOptions{}
		}
		if arg == "--exclude" && i+1 < len(args) {
			opts.ExcludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--include" && i+1 < len(args) {
			opts.IncludeDirs = strings.Split(args[i+1], ",")
		}
		if arg == "--local" {
			useLocal = true
		}
	}

	fmt.Printf("🔒 Security scan: %s\n", repoPath)

	summary, err := scanner.ScanRepoWithOptions(repoPath, opts)
	if err != nil {
		fmt.Printf("❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("✅ Loaded %d files (%d lines) across %d languages\n\n",
		summary.TotalFiles, summary.TotalLines, len(summary.Languages))

	absPath, _ := filepath.Abs(repoPath)
	outDir := filepath.Join(absPath, "security-report")

	fmt.Printf("📝 Generating security report → %s/\n\n", outDir)

	gen := security.NewGenerator(summary, outDir, useLocal)
	if err := gen.Generate(); err != nil {
		fmt.Printf("\n❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\n🎉 Security report complete!\n")
	fmt.Printf("   🎯 security-report/threat-model.md          — STRIDE threat model\n")
	fmt.Printf("   🐛 security-report/vulnerabilities.md       — Vulnerability findings\n")
	fmt.Printf("   🔑 security-report/secrets-audit.md         — Secrets & data exposure\n")
	fmt.Printf("   🛡️  security-report/security-architecture.md — Security design review\n")
	fmt.Printf("   🌐 security-report/index.html               — Interactive report\n")
	fmt.Printf("\n   Open in browser: file://%s/index.html\n", outDir)
}

func cmdScan(repoPath string) {
	fmt.Printf("⏳ Scanning %s...\n", repoPath)

	summary, err := scanner.ScanRepo(repoPath)
	if err != nil {
		fmt.Printf("❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("\n📁 %s\n", summary.Name)
	fmt.Printf("   Files: %d\n", summary.TotalFiles)
	fmt.Printf("   Lines: %d\n", summary.TotalLines)
	fmt.Printf("   Languages:\n")
	for lang, count := range summary.Languages {
		fmt.Printf("     %s: %d files\n", lang, count)
	}

	fmt.Printf("\n📄 File tree:\n")
	for _, f := range summary.Files {
		fmt.Printf("   %s (%d lines)\n", f.RelPath, f.LineCount)
	}
}

func cmdAsk(repoPath string, args []string) {
	// Parse flags
	var singleQuestion string
	useLocal := false
	for i, arg := range args {
		if arg == "-q" && i+1 < len(args) {
			singleQuestion = args[i+1]
		}
		if arg == "--local" {
			useLocal = true
		}
	}

	fmt.Printf("⏳ Scanning %s...\n", repoPath)

	summary, err := scanner.ScanRepo(repoPath)
	if err != nil {
		fmt.Printf("❌ Error: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("✅ Loaded %d files (%d lines) across %d languages\n",
		summary.TotalFiles, summary.TotalLines, len(summary.Languages))

	// Build context
	context := scanner.BuildContext(summary, maxContextChars)

	var client *ollama.Client
	if useLocal {
		client = ollama.NewLocalClient()
	} else {
		client = ollama.NewClient()
	}

	// System message with repo context
	messages := []ollama.Message{
		{
			Role:    "system",
			Content: systemPrompt + "\n\n" + context,
		},
	}

	if singleQuestion != "" {
		// Single question mode
		fmt.Printf("\n🤔 Thinking...\n\n")
		messages = append(messages, ollama.Message{
			Role:    "user",
			Content: singleQuestion,
		})

		_, err := client.ChatStream(messages, func(chunk string) {
			fmt.Print(chunk)
		})
		fmt.Println()
		if err != nil {
			fmt.Printf("\n❌ Error: %v\n", err)
			os.Exit(1)
		}
		return
	}

	// Interactive mode
	fmt.Printf("\n💬 Ask anything about %s\n", summary.Name)
	fmt.Printf("   Type your question and press Enter. Type 'quit' to exit.\n\n")

	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Print("→ ")
		input, err := reader.ReadString('\n')
		if err != nil {
			break
		}

		input = strings.TrimSpace(input)
		if input == "" {
			continue
		}
		if input == "quit" || input == "exit" || input == "q" {
			fmt.Println("Bye!")
			break
		}

		messages = append(messages, ollama.Message{
			Role:    "user",
			Content: input,
		})

		fmt.Println()
		response, err := client.ChatStream(messages, func(chunk string) {
			fmt.Print(chunk)
		})
		fmt.Println("\n")

		if err != nil {
			fmt.Printf("❌ Error: %v\n\n", err)
			continue
		}

		// Add assistant response to history
		messages = append(messages, ollama.Message{
			Role:    "assistant",
			Content: response,
		})

		// Keep history manageable (system + last 20 messages)
		if len(messages) > 21 {
			messages = append(messages[:1], messages[len(messages)-20:]...)
		}
	}
}