inou/portal/mcp_tools.go

package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"inou/lib"
)

// MCP Tool Implementations
// All tools go through lib directly with RBAC enforcement.

// --- Data query tools: all go through lib with RBAC ---

func mcpListDossiers(accessorID string) (string, error) {
	rows, err := lib.DossierQuery(accessorID)
	if err != nil {
		return "", err
	}

	// Group by dossier_id
	type dossierWithCategories struct {
		ID         string   `json:"id"`
		Name       string   `json:"name"`
		Email      string   `json:"email,omitempty"`
		DOB        string   `json:"date_of_birth,omitempty"`
		Sex        string   `json:"sex,omitempty"`
		Categories []string `json:"categories"`
	}

	dossierMap := make(map[string]*dossierWithCategories)
	for _, r := range rows {
		if _, exists := dossierMap[r.DossierID]; !exists {
			dossierMap[r.DossierID] = &dossierWithCategories{
				ID:         r.DossierID,
				Name:       r.Name,
				Email:      r.Email,
				DOB:        r.DateOfBirth,
				Sex:        lib.SexTranslate(r.Sex, "en"),
				Categories: []string{},
			}
		}

		if r.EntryCount > 0 {
			dossierMap[r.DossierID].Categories = append(dossierMap[r.DossierID].Categories, lib.CategoryName(r.Category))
		}
	}

	// Convert map to array
	var result []*dossierWithCategories
	for _, d := range dossierMap {
		result = append(result, d)
	}

	pretty, _ := json.MarshalIndent(result, "", "  ")
	return string(pretty), nil
}

func mcpQueryEntries(accessorID, dossier, category, typ, searchKey, parent, from, to string, limit int) (string, error) {
	cat := -1 // any category
	if category != "" {
		cat = lib.CategoryFromString[category]
	}
	filter := &lib.EntryFilter{DossierID: dossier}
	if typ != "" {
		filter.Type = typ
	}
	if searchKey != "" {
		filter.SearchKey = searchKey
	}
	if from != "" {
		filter.FromDate, _ = strconv.ParseInt(from, 10, 64)
	}
	if to != "" {
		filter.ToDate, _ = strconv.ParseInt(to, 10, 64)
	}
	if limit > 0 {
		filter.Limit = limit
	}

	entries, err := lib.EntryList(accessorID, parent, cat, filter)
	if err != nil {
		return "", err
	}
	return formatEntries(entries), nil
}

func mcpGetCategories(dossier, accessorID string) (string, error) {
	ctx := &lib.AccessContext{AccessorID: accessorID}
	result, err := lib.EntryCategoryCounts(ctx, dossier)
	if err != nil {
		return "", err
	}
	pretty, _ := json.MarshalIndent(result, "", "  ")
	return string(pretty), nil
}

// formatEntries converts entries to the standard MCP response format.
func formatEntries(entries []*lib.Entry) string {
	var result []map[string]any

	// Lazy-load dossier info for lab reference lookups
	var dossierSex string
	var dossierDOB int64
	var dossierLoaded bool

	for _, e := range entries {
		entry := map[string]any{
			"id":        e.EntryID,
			"parent_id": e.ParentID,
			"category":  lib.CategoryName(e.Category),
			"type":      e.Type,
			"value":     e.Value,
			"summary":   e.Summary,
			"ordinal":   e.Ordinal,
			"timestamp": e.Timestamp,
		}
		if e.SearchKey != "" {
			entry["search_key"] = e.SearchKey
		}
		if e.Data != "" {
			var d map[string]any
			if json.Unmarshal([]byte(e.Data), &d) == nil {
				entry["data"] = d
			}
		}
		// Enrich lab result entries with test name, reference range, and flag
		if e.Category == lib.CategoryLab && e.SearchKey != "" {
			test, err := lib.LabTestGet(e.SearchKey)
			if err == nil && test != nil {
				entry["test_name"] = test.Name

				// Load dossier sex/DOB once
				if !dossierLoaded {
					dossierLoaded = true
					if d, err := lib.DossierGet("", e.DossierID); err == nil && d != nil {
						switch d.Sex {
						case 1:
							dossierSex = "M"
						case 2:
							dossierSex = "F"
						}
						dossierDOB = d.DOB.Unix()
					}
				}

				// Look up reference range for this test at the patient's age at time of lab
				ts := e.Timestamp
				if ts == 0 {
					ts = int64(e.Ordinal) // fallback
				}
				if dossierDOB > 0 && ts > 0 {
					ageDays := lib.AgeDays(dossierDOB, ts)
					if ref, err := lib.LabRefLookup(e.SearchKey, dossierSex, ageDays); err == nil && ref != nil {
						siFactor := float64(test.SIFactor) / lib.LabScale
						if siFactor > 0 {
							low := lib.FromLabScale(ref.RefLow) / siFactor
							high := lib.FromLabScale(ref.RefHigh) / siFactor
							entry["ref_low"] = low
							entry["ref_high"] = high

							// Compute flag from numeric value
							if numVal, ok := entry["value"].(string); ok {
								if v, err := strconv.ParseFloat(numVal, 64); err == nil {
									if ref.RefLow >= 0 && v < low {
										entry["flag"] = "L"
									} else if ref.RefHigh >= 0 && v > high {
										entry["flag"] = "H"
									}
								}
							}
						}
					}
				}
			}
		}
		switch e.Type {
		case "root":
			entry["hint"] = "Use list_entries with parent=" + e.EntryID + " to list studies"
		case "study":
			entry["hint"] = "Use list_entries with parent=" + e.EntryID + " to list series"
		case "series":
			entry["hint"] = "Use fetch_contact_sheet with series=" + e.EntryID + " to browse slices, then fetch_image with the slice ID"
		}
		result = append(result, entry)
	}
	pretty, _ := json.MarshalIndent(result, "", "  ")
	return string(pretty)
}

// --- Image tools: direct through lib ---

func mcpFetchImage(accessorID, dossier, slice string, wc, ww float64) (map[string]interface{}, error) {
	opts := &lib.ImageOpts{WC: wc, WW: ww}
	body, err := lib.RenderImage(accessorID, slice, opts, 2000)
	if err != nil {
		return nil, err
	}
	b64 := base64.StdEncoding.EncodeToString(body)
	return mcpImageContent(b64, "image/webp", fmt.Sprintf("Slice %s (%d bytes)", slice[:8], len(body))), nil
}

func mcpFetchContactSheet(accessorID, dossier, series string, wc, ww float64) (map[string]interface{}, error) {
	body, err := lib.RenderContactSheet(accessorID, series, wc, ww)
	if err != nil {
		return nil, err
	}
	b64 := base64.StdEncoding.EncodeToString(body)
	return mcpImageContent(b64, "image/webp", fmt.Sprintf("Contact sheet %s (%d bytes)", series[:8], len(body))), nil
}

// --- Document fetch: returns extracted text + metadata from Data field ---
// mcpFetchDocument returns a full MCP content map.
// format: "original" = base64 PDF, "markdown" = formatted text, "translation" = translated text

func mcpFetchDocument(accessorID, dossier, entryID, format string) (map[string]interface{}, error) {
	// Use EntryGet (by ID only) — EntryRead with Category=0 default would exclude non-profile entries.
	e, err := lib.EntryGet(&lib.AccessContext{AccessorID: accessorID}, entryID)
	if err != nil {
		return nil, err
	}
	if e == nil {
		return nil, fmt.Errorf("document not found")
	}
	// Verify the entry belongs to the requested dossier.
	if e.DossierID != dossier {
		return nil, fmt.Errorf("document not found")
	}

	// Parse the Data field (populated by doc-processor).
	var data map[string]interface{}
	if e.Data != "" {
		_ = json.Unmarshal([]byte(e.Data), &data)
	}

	if format == "" {
		format = "original"
	}

	switch format {
	case "markdown":
		text := docToMarkdown(e, data)
		return mcpTextContent(text), nil
	case "translation":
		text, err := docToTranslation(e, data)
		if err != nil {
			return nil, err
		}
		return mcpTextContent(text), nil
	default: // "original" — return base64-encoded PDF
		return docToOriginalPDF(e, data)
	}
}

// docToOriginalPDF decrypts the source PDF and returns it as base64 MCP content.
func docToOriginalPDF(e *lib.Entry, data map[string]interface{}) (map[string]interface{}, error) {
	sourceUpload, _ := data["source_upload"].(string)
	if sourceUpload == "" {
		return nil, fmt.Errorf("no PDF available for this document")
	}

	uploadEntry, err := lib.EntryGet(nil, sourceUpload)
	if err != nil || uploadEntry == nil {
		return nil, fmt.Errorf("upload entry not found")
	}

	var uploadData struct {
		Path string `json:"path"`
	}
	if err := json.Unmarshal([]byte(uploadEntry.Data), &uploadData); err != nil || uploadData.Path == "" {
		return nil, fmt.Errorf("no file path in upload entry")
	}

	pdfBytes, err := lib.DecryptFile(uploadData.Path)
	if err != nil {
		return nil, fmt.Errorf("decrypt failed: %w", err)
	}

	b64 := base64.StdEncoding.EncodeToString(pdfBytes)
	summary := e.Summary
	if summary == "" {
		summary = "document"
	}

	return map[string]interface{}{
		"content": []map[string]interface{}{
			{
				"type":     "resource",
				"resource": map[string]interface{}{
					"uri":      "data:application/pdf;base64," + b64,
					"mimeType": "application/pdf",
					"text":     summary,
				},
			},
		},
	}, nil
}

// docToMarkdown returns the pre-rendered markdown stored by doc-processor.
func docToMarkdown(e *lib.Entry, data map[string]interface{}) string {
	if md, ok := data["markdown"].(string); ok && md != "" {
		return md
	}
	// Fallback: summary only
	return e.Summary
}

// docToTranslation returns the pre-translated markdown if available,
// otherwise translates the markdown field on-the-fly via Claude.
func docToTranslation(e *lib.Entry, data map[string]interface{}) (string, error) {
	// Use pre-translated version if already stored by doc-processor.
	if tr, ok := data["markdown_translated"].(string); ok && tr != "" {
		return tr, nil
	}

	// Fall back to on-the-fly translation.
	src, _ := data["markdown"].(string)
	if src == "" {
		src = e.Summary
	}
	if src == "" {
		return "", fmt.Errorf("no text content to translate")
	}
	if lib.AnthropicKey == "" {
		return "", fmt.Errorf("translation unavailable: no Anthropic API key configured")
	}

	prompt := "Translate the following medical document (markdown format) to English. Preserve all markdown formatting, medical terminology, values, and structure. Output only the translated markdown, no explanation.\n\n" + src

	reqBody, _ := json.Marshal(map[string]interface{}{
		"model":      "claude-haiku-4-5",
		"max_tokens": 4096,
		"messages": []map[string]interface{}{
			{"role": "user", "content": prompt},
		},
	})

	req, err := http.NewRequest("POST", "https://api.anthropic.com/v1/messages", strings.NewReader(string(reqBody)))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("x-api-key", lib.AnthropicKey)
	req.Header.Set("anthropic-version", "2023-06-01")

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

	var result struct {
		Content []struct {
			Text string `json:"text"`
		} `json:"content"`
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("translation API error: %s", result.Error.Message)
	}
	if len(result.Content) == 0 {
		return "", fmt.Errorf("empty translation response")
	}
	return result.Content[0].Text, nil
}