markdown-hub/internal/api/ai.go

package api

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

	"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"`
	}
	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, 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 and return the COMPLETE updated document. " +
			"Do not add explanations or wrap in code fences. Return raw markdown 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. Answer concisely in markdown."
		userMsg = "Document:\n\n" + req.Content + "\n\nQuestion: " + req.Message
	}

	// Stream SSE response
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeJSON(w, 500, map[string]string{"error": "streaming not supported"})
		return
	}

	payload := map[string]interface{}{
		"model":       aiModel,
		"messages":    []map[string]string{{"role": "system", "content": systemPrompt}, {"role": "user", "content": userMsg}},
		"temperature": 0.3,
		"stream":      true,
	}
	body, _ := json.Marshal(payload)
	url := strings.TrimRight(aiEndpoint, "/") + "/chat/completions"
	aiReq, _ := http.NewRequest("POST", url, strings.NewReader(string(body)))
	aiReq.Header.Set("Content-Type", "application/json")
	if aiKey != "" {
		aiReq.Header.Set("Authorization", "Bearer "+aiKey)
	}

	client := &http.Client{Timeout: 120 * time.Second}
	resp, err := client.Do(aiReq)
	if err != nil {
		fmt.Fprintf(w, "data: {\"error\":\"AI unreachable\"}\n\n")
		flusher.Flush()
		return
	}
	defer resp.Body.Close()

	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		data := line[6:]
		if data == "[DONE]" {
			fmt.Fprintf(w, "data: [DONE]\n\n")
			flusher.Flush()
			break
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if json.Unmarshal([]byte(data), &chunk) == nil && len(chunk.Choices) > 0 {
			token := chunk.Choices[0].Delta.Content
			if token != "" {
				tokenJSON, _ := json.Marshal(token)
				fmt.Fprintf(w, "data: %s\n\n", tokenJSON)
				flusher.Flush()
			}
		}
	}
}

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},
		},
		"temperature": 0.3,
	}
	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)
	}

	client := &http.Client{Timeout: 120 * time.Second}
	resp, err := client.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
}