package main

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

	"inou/lib"
)

// MCP HTTP Server Implementation
// Implements Streamable HTTP transport for Anthropic's Connectors Directory
// See: https://modelcontextprotocol.io/specification/draft/basic/authorization

const (
	mcpProtocolVersion = "2025-06-18"
	mcpServerName      = "inou-health"
	mcpServerVersion   = "1.0.0"
)

// OAuth 2.0 Protected Resource Metadata (RFC9728)
// GET /.well-known/oauth-protected-resource
func handleOAuthProtectedResource(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Determine the base URL from the request
	scheme := "https"
	if r.TLS == nil && r.Header.Get("X-Forwarded-Proto") != "https" {
		scheme = "http"
	}
	host := r.Host
	baseURL := fmt.Sprintf("%s://%s", scheme, host)

	metadata := map[string]interface{}{
		"resource":              baseURL + "/mcp",
		"authorization_servers": []string{baseURL},
		"scopes_supported":      []string{"read"},
		"bearer_methods_supported": []string{"header"},
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "max-age=3600")
	json.NewEncoder(w).Encode(metadata)
}

// OAuth 2.0 Authorization Server Metadata (RFC8414)
// GET /.well-known/oauth-authorization-server
func handleOAuthAuthorizationServer(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	scheme := "https"
	if r.TLS == nil && r.Header.Get("X-Forwarded-Proto") != "https" {
		scheme = "http"
	}
	host := r.Host
	baseURL := fmt.Sprintf("%s://%s", scheme, host)

	metadata := map[string]interface{}{
		"issuer":                                baseURL,
		"authorization_endpoint":                baseURL + "/oauth/authorize",
		"token_endpoint":                        baseURL + "/oauth/token",
		"userinfo_endpoint":                     baseURL + "/oauth/userinfo",
		"revocation_endpoint":                   baseURL + "/oauth/revoke",
		"registration_endpoint":                 baseURL + "/register",
		"response_types_supported":              []string{"code"},
		"grant_types_supported":                 []string{"authorization_code", "refresh_token"},
		"token_endpoint_auth_methods_supported": []string{"none", "client_secret_post"},
		"code_challenge_methods_supported":      []string{"S256"},
		"scopes_supported":                      []string{"read"},
	}

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Cache-Control", "max-age=3600")
	json.NewEncoder(w).Encode(metadata)
}

// Dynamic Client Registration (RFC 7591)
// POST /register - Allows clients to register themselves dynamically
func handleDynamicClientRegistration(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Parse registration request
	var req struct {
		RedirectURIs            []string `json:"redirect_uris"`
		ClientName              string   `json:"client_name"`
		TokenEndpointAuthMethod string   `json:"token_endpoint_auth_method"`
		GrantTypes              []string `json:"grant_types"`
		ResponseTypes           []string `json:"response_types"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(map[string]string{
			"error":             "invalid_client_metadata",
			"error_description": "Invalid JSON body",
		})
		return
	}

	// Validate redirect URIs
	if len(req.RedirectURIs) == 0 {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(map[string]string{
			"error":             "invalid_redirect_uri",
			"error_description": "At least one redirect_uri is required",
		})
		return
	}

	// Create client name if not provided
	clientName := req.ClientName
	if clientName == "" {
		clientName = "Dynamic Client"
	}

	// Create the client
	client, secret, err := lib.OAuthClientCreate(clientName, req.RedirectURIs)
	if err != nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusInternalServerError)
		json.NewEncoder(w).Encode(map[string]string{
			"error":             "server_error",
			"error_description": "Failed to create client",
		})
		return
	}

	// Return registration response
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(map[string]interface{}{
		"client_id":                  client.ClientID,
		"client_secret":              secret,
		"client_name":                clientName,
		"redirect_uris":              req.RedirectURIs,
		"token_endpoint_auth_method": "client_secret_post",
		"grant_types":                []string{"authorization_code", "refresh_token"},
		"response_types":             []string{"code"},
	})
}

// MCP JSON-RPC types
type mcpRequest struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      interface{}     `json:"id"`
	Method  string          `json:"method"`
	Params  json.RawMessage `json:"params,omitempty"`
}

type mcpResponse struct {
	JSONRPC string      `json:"jsonrpc"`
	ID      interface{} `json:"id"`
	Result  interface{} `json:"result,omitempty"`
	Error   *mcpError   `json:"error,omitempty"`
}

type mcpError struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
}

// MCP HTTP endpoint
// POST /mcp - Streamable HTTP transport
func handleMCP(w http.ResponseWriter, r *http.Request) {
	// FIXED(review-2026-02-28): Use origin allowlist instead of wildcard
	origin := r.Header.Get("Origin")
	if corsAllowedOrigins[origin] {
		w.Header().Set("Access-Control-Allow-Origin", origin)
		w.Header().Set("Access-Control-Allow-Credentials", "true")
	}
	w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

	if r.Method == "OPTIONS" {
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Check authorization
	var dossierID string

	// Local LAN backdoor — skip OAuth for agents on 192.168.0.0/22
	remoteIP := strings.Split(r.RemoteAddr, ":")[0]
	if strings.HasPrefix(remoteIP, "192.168.0.") || strings.HasPrefix(remoteIP, "192.168.1.") || strings.HasPrefix(remoteIP, "192.168.2.") || strings.HasPrefix(remoteIP, "192.168.3.") {
		dossierID = "6e4e8192881a7494"
		fmt.Printf("[MCP] Local gateway access as dossier: %s\n", dossierID)
	} else {
		auth := r.Header.Get("Authorization")
		scheme := "https"
		if r.TLS == nil && r.Header.Get("X-Forwarded-Proto") != "https" {
			scheme = "http"
		}
		baseURL := fmt.Sprintf("%s://%s", scheme, r.Host)

		if auth == "" || !strings.HasPrefix(auth, "Bearer ") {
			w.Header().Set("WWW-Authenticate", fmt.Sprintf(
				`Bearer resource_metadata="%s/.well-known/oauth-protected-resource", scope="read"`,
				baseURL,
			))
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		tokenStr := strings.TrimPrefix(auth, "Bearer ")
		token, err := lib.TokenParse(tokenStr)
		if err != nil {
			w.Header().Set("WWW-Authenticate", fmt.Sprintf(
				`Bearer resource_metadata="%s/.well-known/oauth-protected-resource", error="invalid_token"`,
				baseURL,
			))
			http.Error(w, "Invalid token", http.StatusUnauthorized)
			return
		}

		fmt.Printf("[MCP] OAuth token authenticated as dossier: %s\n", token.DossierID)

		if d, _ := lib.DossierGet(lib.SystemAccessorID, token.DossierID); d == nil {
			w.Header().Set("WWW-Authenticate", fmt.Sprintf(
				`Bearer resource_metadata="%s/.well-known/oauth-protected-resource", error="invalid_token", error_description="Dossier not found"`,
				baseURL,
			))
			http.Error(w, "Invalid token: dossier not found", http.StatusUnauthorized)
			return
		}

		dossierID = token.DossierID
	}
	body, err := io.ReadAll(r.Body)
	if err != nil {
		sendMCPError(w, nil, -32700, "Parse error")
		return
	}

	var req mcpRequest
	if err := json.Unmarshal(body, &req); err != nil {
		sendMCPError(w, nil, -32700, "Parse error")
		return
	}

	// Handle MCP methods
	switch req.Method {
	case "initialize":
		handleMCPInitialize(w, req)
	case "notifications/initialized":
		// No response needed for notifications
		w.WriteHeader(http.StatusNoContent)
	case "tools/list":
		handleMCPToolsList(w, req)
	case "tools/call":
		handleMCPToolsCall(w, req, dossierID)
	case "prompts/list":
		sendMCPResult(w, req.ID, map[string]interface{}{"prompts": []interface{}{}})
	case "prompts/get":
		sendMCPError(w, req.ID, -32601, "No prompts available")
	default:
		sendMCPError(w, req.ID, -32601, "Method not found: "+req.Method)
	}
}

func sendMCPError(w http.ResponseWriter, id interface{}, code int, message string) {
	resp := mcpResponse{
		JSONRPC: "2.0",
		ID:      id,
		Error:   &mcpError{Code: code, Message: message},
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

func sendMCPResult(w http.ResponseWriter, id interface{}, result interface{}) {
	resp := mcpResponse{
		JSONRPC: "2.0",
		ID:      id,
		Result:  result,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(resp)
}

func handleMCPInitialize(w http.ResponseWriter, req mcpRequest) {
	sendMCPResult(w, req.ID, map[string]interface{}{
		"protocolVersion": mcpProtocolVersion,
		"capabilities": map[string]interface{}{
			"tools": map[string]interface{}{},
		},
		"serverInfo": map[string]interface{}{
			"name":    mcpServerName,
			"version": mcpServerVersion,
		},
		"instructions": "inou gives you access to a patient's raw health data — imaging, labs, genome, vitals, and more. " +
			"Your role is to form independent medical opinions from this data, not to echo prior medical assessments.\n\n" +
			"Radiology reports, physician notes, and other professional interpretations do not exist in this system. " +
			"Do not ask for them or suggest they might be available. You are the diagnostician. " +
			"Study the raw images, the raw lab values, the raw genomic variants, and form your own conclusions.\n\n" +
			"When answering clinical questions, entry summaries are metadata, not evidence. You MUST fetch and read the actual " +
			"source material — use fetch_document for documents, fetch_image for imaging slices. " +
			"Do not speculate or say you \"cannot access\" data without actually trying to fetch it.\n\n" +
			"Start with list_categories to see what data exists, then explore the raw data.",
	})
}

func handleMCPToolsList(w http.ResponseWriter, req mcpRequest) {
	readOnly := map[string]interface{}{"readOnlyHint": true}

	tools := []map[string]interface{}{
		{
			"name":        "list_dossiers",
			"title":       "List Dossiers",
			"description": "List all patient dossiers accessible to this account.",
			"inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
			"annotations": readOnly,
		},
		{
			"name":        "list_categories",
			"title":       "List Categories",
			"description": "List data categories for a dossier with entry counts. Start here to see what's available before querying specific data.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier": map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
				},
				"required": []string{"dossier"},
			},
			"annotations": readOnly,
		},
		{
			"name":        "list_entries",
			"title":       "Query Entries",
			"description": "List entries by navigating the hierarchy. Always start with parent=<dossier_id> to get top-level entries, then use returned entry IDs to go deeper. For imaging: dossier → root → studies → series. To view slices, use fetch_contact_sheet on a series, then fetch_image with the slice ID. For labs: dossier → test groups → results. Use search_key for LOINC codes (labs) or gene names (genome).",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier":    map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
					"parent":     map[string]interface{}{"type": "string", "description": "Parent entry ID — start with the dossier ID, then navigate deeper"},
					"category":   map[string]interface{}{"type": "string", "description": "Category name (use list_categories to discover)"},
					"type":       map[string]interface{}{"type": "string", "description": "Entry type within category"},
					"search_key": map[string]interface{}{"type": "string", "description": "LOINC code for labs, gene name for genome"},
					"from":       map[string]interface{}{"type": "string", "description": "Timestamp start (Unix seconds)"},
					"to":         map[string]interface{}{"type": "string", "description": "Timestamp end (Unix seconds)"},
					"limit":      map[string]interface{}{"type": "number", "description": "Maximum results"},
				},
				"required": []string{"dossier", "parent"},
			},
			"annotations": readOnly,
		},
		{
			"name":        "fetch_image",
			"title":       "Fetch Image",
			"description": "Fetch slice image as base64 PNG. Optionally set window/level.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier": map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
					"slice":   map[string]interface{}{"type": "string", "description": "Slice ID (16-char hex)"},
					"wc":      map[string]interface{}{"type": "number", "description": "Window center"},
					"ww":      map[string]interface{}{"type": "number", "description": "Window width"},
				},
				"required": []string{"dossier", "slice"},
			},
			"annotations": readOnly,
		},
		{
			"name":        "fetch_contact_sheet",
			"title":       "Fetch Contact Sheet",
			"description": "Fetch contact sheet (thumbnail grid) for NAVIGATION ONLY. Use to identify slices, then fetch at full resolution. NEVER diagnose from thumbnails.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier": map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
					"series":  map[string]interface{}{"type": "string", "description": "Series ID (16-char hex)"},
					"wc":      map[string]interface{}{"type": "number", "description": "Window center"},
					"ww":      map[string]interface{}{"type": "number", "description": "Window width"},
				},
				"required": []string{"dossier", "series"},
			},
			"annotations": readOnly,
		},
		{
			"name":        "fetch_document",
			"title":       "Fetch Document",
			"description": "Fetch full document content including extracted text, findings, and metadata. Use after finding documents via list_entries.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier":  map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
					"entry_id": map[string]interface{}{"type": "string", "description": "Document entry ID (16-char hex)"},
					"format":   map[string]interface{}{"type": "string", "description": "Output format: 'original' (default, raw JSON), 'markdown' (formatted), 'translation' (English translation via AI)"},
				},
				"required": []string{"dossier", "entry_id"},
			},
			"annotations": readOnly,
		},
		{
			"name":        "get_version",
			"title":       "Server Version",
			"description": "Get server version info.",
			"inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
			"annotations": readOnly,
		},
		{
			"name":        "create_entry",
			"title":       "Create Entry",
			"description": "Write a new health entry directly into a dossier. Use for saving AI-generated insights, supplement notes, nutrition logs, or any health observation. Requires write permission on the dossier. Category must be one of the values returned by list_categories (e.g. 'supplement', 'nutrition', 'tracker'). Type is a sub-type within the category (e.g. 'vitamin', 'meal', 'note'). Value is a short human-readable label. Summary is a one-line description. Data is optional structured JSON for additional fields.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dossier":   map[string]interface{}{"type": "string", "description": "Dossier ID (16-char hex)"},
					"category":  map[string]interface{}{"type": "string", "description": "Category name (e.g. 'supplement', 'nutrition', 'tracker')"},
					"type":      map[string]interface{}{"type": "string", "description": "Entry type within category (e.g. 'vitamin', 'meal', 'note')"},
					"value":     map[string]interface{}{"type": "string", "description": "Short label or title for the entry"},
					"summary":   map[string]interface{}{"type": "string", "description": "One-line description or AI-generated insight"},
					"data":      map[string]interface{}{"type": "string", "description": "Optional JSON string with additional structured fields"},
					"parent":    map[string]interface{}{"type": "string", "description": "Optional parent entry ID for hierarchical entries"},
					"timestamp": map[string]interface{}{"type": "number", "description": "Optional Unix timestamp (defaults to now)"},
				},
				"required": []string{"dossier", "category", "value"},
			},
		},
	}

	sendMCPResult(w, req.ID, map[string]interface{}{"tools": tools})
}

func handleMCPToolsCall(w http.ResponseWriter, req mcpRequest, dossierID string) {
	var params struct {
		Name      string                 `json:"name"`
		Arguments map[string]interface{} `json:"arguments"`
	}
	if err := json.Unmarshal(req.Params, &params); err != nil {
		sendMCPError(w, req.ID, -32602, "Invalid params")
		return
	}

	switch params.Name {
	case "list_dossiers":
		result, err := mcpListDossiers(dossierID)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, mcpTextContent(result))

	case "list_categories":
		dossier, _ := params.Arguments["dossier"].(string)
		if dossier == "" {
			sendMCPError(w, req.ID, -32602, "dossier required")
			return
		}
		result, err := mcpGetCategories(dossier, dossierID)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, mcpTextContent(result))

	case "list_entries":
		dossier, _ := params.Arguments["dossier"].(string)
		if dossier == "" {
			sendMCPError(w, req.ID, -32602, "dossier required")
			return
		}
		category, _ := params.Arguments["category"].(string)
		typ, _ := params.Arguments["type"].(string)
		searchKey, _ := params.Arguments["search_key"].(string)
		parent, _ := params.Arguments["parent"].(string)
		if parent == "" && searchKey == "" {
			sendMCPResult(w, req.ID, mcpTextContent("ERROR: parent is required. Start with parent="+dossier+" (the dossier ID) to list top-level entries, then use returned entry IDs to navigate deeper."))
			return
		}
		// When parent is the dossier ID, search all top-level entries (ParentID="" or ParentID=dossierID)
		if parent == dossier {
			parent = ""
		}
		from, _ := params.Arguments["from"].(string)
		to, _ := params.Arguments["to"].(string)
		limit, _ := params.Arguments["limit"].(float64)
		result, err := mcpQueryEntries(dossierID, dossier, category, typ, searchKey, parent, from, to, int(limit))
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, mcpTextContent(result))

	case "fetch_image":
		dossier, _ := params.Arguments["dossier"].(string)
		slice, _ := params.Arguments["slice"].(string)
		if dossier == "" || slice == "" {
			sendMCPError(w, req.ID, -32602, "dossier and slice required")
			return
		}
		wc, _ := params.Arguments["wc"].(float64)
		ww, _ := params.Arguments["ww"].(float64)
		result, err := mcpFetchImage(dossierID, dossier, slice, wc, ww)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, result)

	case "fetch_contact_sheet":
		dossier, _ := params.Arguments["dossier"].(string)
		series, _ := params.Arguments["series"].(string)
		if dossier == "" || series == "" {
			sendMCPError(w, req.ID, -32602, "dossier and series required")
			return
		}
		wc, _ := params.Arguments["wc"].(float64)
		ww, _ := params.Arguments["ww"].(float64)
		result, err := mcpFetchContactSheet(dossierID, dossier, series, wc, ww)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, result)

	case "fetch_document":
		dossier, _ := params.Arguments["dossier"].(string)
		entryID, _ := params.Arguments["entry_id"].(string)
		format, _ := params.Arguments["format"].(string)
		if dossier == "" || entryID == "" {
			sendMCPError(w, req.ID, -32602, "dossier and entry_id required")
			return
		}
		result, err := mcpFetchDocument(dossierID, dossier, entryID, format)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, result)

	case "get_version":
		sendMCPResult(w, req.ID, mcpTextContent(fmt.Sprintf("Server: %s v%s", mcpServerName, mcpServerVersion)))

	case "create_entry":
		dossier, _ := params.Arguments["dossier"].(string)
		category, _ := params.Arguments["category"].(string)
		typ, _ := params.Arguments["type"].(string)
		value, _ := params.Arguments["value"].(string)
		summary, _ := params.Arguments["summary"].(string)
		data, _ := params.Arguments["data"].(string)
		parent, _ := params.Arguments["parent"].(string)
		var ts int64
		if tsRaw, ok := params.Arguments["timestamp"]; ok {
			switch v := tsRaw.(type) {
			case float64:
				ts = int64(v)
			case int64:
				ts = v
			}
		}
		result, err := mcpCreateEntry(dossierID, dossier, category, typ, value, summary, data, parent, ts)
		if err != nil {
			sendMCPError(w, req.ID, -32000, err.Error())
			return
		}
		sendMCPResult(w, req.ID, mcpTextContent(result))

	default:
		sendMCPError(w, req.ID, -32601, "Unknown tool: "+params.Name)
	}
}

func mcpTextContent(text string) map[string]interface{} {
	return map[string]interface{}{
		"content": []map[string]interface{}{
			{"type": "text", "text": text},
		},
	}
}

func mcpImageContent(base64Data, mimeType, text string) map[string]interface{} {
	return map[string]interface{}{
		"content": []map[string]interface{}{
			{"type": "image", "data": base64Data, "mimeType": mimeType},
			{"type": "text", "text": text},
		},
	}
}

// RegisterMCPRoutes registers MCP HTTP endpoints
func RegisterMCPRoutes(mux *http.ServeMux) {
	mux.HandleFunc("/.well-known/oauth-protected-resource", handleOAuthProtectedResource)
	mux.HandleFunc("/.well-known/oauth-protected-resource/mcp", handleOAuthProtectedResource)
	mux.HandleFunc("/.well-known/oauth-authorization-server", handleOAuthAuthorizationServer)
	mux.HandleFunc("/register", handleDynamicClientRegistration)
	mux.HandleFunc("/mcp", handleMCP)
}