markdown-hub/internal/api/ai.go

package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"

	"markdownhub/internal/files"
)

func (s *Server) handleAIVerify(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Path string `json:"path"`
	}
	if err := decodeBody(r, &req); err != nil || req.Path == "" {
		writeJSON(w, 400, map[string]string{"error": "path required"})
		return
	}

	userID := getUserID(r)
	content, err := files.ReadFile(s.dataDir, userID, req.Path)
	if err != nil {
		writeJSON(w, 404, map[string]string{"error": "file not found"})
		return
	}

	aiEndpoint := os.Getenv("MH_AI_ENDPOINT")
	aiKey := os.Getenv("MH_AI_API_KEY")
	aiModel := os.Getenv("MH_AI_MODEL")
	if aiEndpoint == "" {
		writeJSON(w, 500, map[string]string{"error": "AI endpoint not configured (MH_AI_ENDPOINT)"})
		return
	}
	if aiModel == "" {
		aiModel = "gpt-4"
	}

	// Call LiteLLM-compatible endpoint
	systemPrompt := `You are a technical reviewer. Review the following specification document for:
1. Completeness - are there missing details needed to implement this?
2. Ambiguities - are there unclear requirements?
3. Feasibility - is this technically achievable?
4. Suggestions - any improvements?

Respond with a structured review. End with a clear verdict: READY TO BUILD or NEEDS REVISION.`

	response, err := callLLM(aiEndpoint, aiKey, aiModel, systemPrompt, content)
	if err != nil {
		writeJSON(w, 500, map[string]string{"error": "AI call failed: " + err.Error()})
		return
	}

	ready := strings.Contains(strings.ToUpper(response), "READY TO BUILD")
	writeJSON(w, 200, map[string]interface{}{
		"feedback": response,
		"ready":    ready,
	})
}

func (s *Server) handleAIGenerate(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Path         string `json:"path"`
		Selection    string `json:"selection"`
		Action       string `json:"action"`
		OutputFolder string `json:"output_folder"`
	}
	if err := decodeBody(r, &req); err != nil {
		writeJSON(w, 400, map[string]string{"error": "invalid request"})
		return
	}

	userID := getUserID(r)
	var inputText string
	if req.Selection != "" {
		inputText = req.Selection
	} else if req.Path != "" {
		content, err := files.ReadFile(s.dataDir, userID, req.Path)
		if err != nil {
			writeJSON(w, 404, map[string]string{"error": "file not found"})
			return
		}
		inputText = content
	}

	aiEndpoint := os.Getenv("MH_AI_ENDPOINT")
	aiKey := os.Getenv("MH_AI_API_KEY")
	aiModel := os.Getenv("MH_AI_MODEL")
	if aiEndpoint == "" {
		writeJSON(w, 500, map[string]string{"error": "AI endpoint not configured"})
		return
	}
	if aiModel == "" {
		aiModel = "gpt-4"
	}

	systemPrompt := "You are a helpful assistant. Respond in markdown."
	switch req.Action {
	case "summarize", "summary":
		systemPrompt = "Summarize the following text concisely in markdown."
	case "prompt":
		systemPrompt = "Generate a detailed AI prompt based on the following specification. The prompt should instruct an AI coding agent to implement the project."
	case "expand":
		systemPrompt = "Expand on the following text with more detail, examples, and explanations. Respond in markdown."
	case "grammar":
		systemPrompt = "Review the following text for grammar and spelling errors. List each error with the correction. Be concise. Format as a markdown list."
	case "spec":
		systemPrompt = `You are a technical reviewer. Review the following specification document for:
1. Completeness - are there missing details needed to implement this?
2. Ambiguities - are there unclear requirements?
3. Feasibility - is this technically achievable?
4. Suggestions - any improvements?
Respond with a structured review. End with a clear verdict: READY TO BUILD or NEEDS REVISION.`
	}

	response, err := callLLM(aiEndpoint, aiKey, aiModel, systemPrompt, inputText)
	if err != nil {
		writeJSON(w, 500, map[string]string{"error": "AI call failed: " + err.Error()})
		return
	}

	// Optionally save to folder
	if req.OutputFolder != "" {
		filename := fmt.Sprintf("%s/%s-output.md", req.OutputFolder, req.Action)
		files.WriteFile(s.dataDir, userID, filename, response)
	}

	writeJSON(w, 200, map[string]string{"result": response})
}

func (s *Server) handleAIChat(w http.ResponseWriter, r *http.Request) {
	var req struct {
		Path    string `json:"path"`
		Content string `json:"content"`
		Message string `json:"message"`
		Mode    string `json:"mode"` // "edit" or "chat"
	}
	if err := decodeBody(r, &req); err != nil || req.Message == "" {
		writeJSON(w, 400, map[string]string{"error": "message required"})
		return
	}

	aiEndpoint := os.Getenv("MH_AI_ENDPOINT")
	aiKey := os.Getenv("MH_AI_API_KEY")
	aiModel := os.Getenv("MH_AI_MODEL")
	if aiEndpoint == "" {
		writeJSON(w, 500, map[string]string{"error": "AI endpoint not configured"})
		return
	}
	if aiModel == "" {
		aiModel = "gpt-4"
	}

	var systemPrompt string
	var userMsg string

	if req.Mode == "edit" {
		systemPrompt = "You are a document editor. The user will give you a markdown document and an instruction. " +
			"Apply the instruction to the document and return ONLY the complete updated document. " +
			"Do not add explanations, comments, or wrap the output in code fences. " +
			"Do not prefix with any markers. Preserve the document's existing style and formatting. " +
			"Return the raw markdown content only."
		userMsg = "Document:\n\n" + req.Content + "\n\nInstruction: " + req.Message
	} else {
		systemPrompt = `You are a helpful writing assistant. The user has a markdown document open and is asking a question about it.
Answer concisely in markdown. Reference the document content when relevant.`
		userMsg = "Document:\n\n" + req.Content + "\n\nQuestion: " + req.Message
	}

	response, err := callLLM(aiEndpoint, aiKey, aiModel, systemPrompt, userMsg)
	if err != nil {
		writeJSON(w, 500, map[string]string{"error": "AI call failed: " + err.Error()})
		return
	}

	if req.Mode == "edit" {
		// Strip markdown code fences if AI wrapped the output
		response = strings.TrimSpace(response)
		if strings.HasPrefix(response, "```") {
			lines := strings.Split(response, "\n")
			if len(lines) > 2 {
				lines = lines[1:] // remove opening fence
				if strings.TrimSpace(lines[len(lines)-1]) == "```" {
					lines = lines[:len(lines)-1] // remove closing fence
				}
				response = strings.Join(lines, "\n")
			}
		}
		writeJSON(w, 200, map[string]string{"content": response})
	} else {
		writeJSON(w, 200, map[string]string{"result": response})
	}
}

func callLLM(endpoint, apiKey, model, systemPrompt, userContent string) (string, error) {
	payload := map[string]interface{}{
		"model": model,
		"messages": []map[string]string{
			{"role": "system", "content": systemPrompt},
			{"role": "user", "content": userContent},
		},
	}
	body, _ := json.Marshal(payload)

	url := strings.TrimRight(endpoint, "/") + "/chat/completions"
	req, err := http.NewRequest("POST", url, strings.NewReader(string(body)))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	if resp.StatusCode != 200 {
		return "", fmt.Errorf("LLM returned %d: %s", resp.StatusCode, string(respBody))
	}

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(respBody, &result); err != nil {
		return "", err
	}
	if len(result.Choices) == 0 {
		return "", fmt.Errorf("no response from LLM")
	}
	return result.Choices[0].Message.Content, nil
}