package main

import (
	"bufio"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"crypto/tls"
	"net/http"
	"net/url"
	"os"
	"runtime"
	"strings"
	"time"
)

var (
	Version   = "dev"
	BuildTime = "unknown"
)

var httpClient = &http.Client{\n	Transport: &http.Transport{\n		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},\n	},\n}

var (
	server string
	token  string
)

var httpClient = &http.Client{\n	Transport: &http.Transport{\n		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},\n	},\n}

func log(format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	fmt.Fprintf(os.Stderr, "[%s] %s\n", time.Now().Format("15:04:05.000"), msg)
}

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"`
}

func respond(id interface{}, result interface{}) {
	resp := MCPResponse{JSONRPC: "2.0", ID: id, Result: result}
	data, _ := json.Marshal(resp)
	fmt.Println(string(data))
}

func respondError(id interface{}, code int, message string) {
	resp := MCPResponse{JSONRPC: "2.0", ID: id, Error: &MCPError{Code: code, Message: message}}
	data, _ := json.Marshal(resp)
	fmt.Println(string(data))
}

func respondContent(id interface{}, text string) {
	respond(id, map[string]interface{}{
		"content": []map[string]interface{}{
			{"type": "text", "text": text},
		},
	})
}

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

func apiGet(endpoint string, params map[string]string) ([]byte, error) {
	u := server + endpoint
	if len(params) > 0 {
		v := url.Values{}
		for k, val := range params {
			v.Set(k, val)
		}
		u += "?" + v.Encode()
	}
	if !strings.Contains(u, "token=") {
		if strings.Contains(u, "?") {
			u += "&token=" + token
		} else {
			u += "?token=" + token
		}
	}
	log("GET %s", u)

	resp, err := http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

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

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("HTTP %d: %s", resp.StatusCode, string(body))
	}

	return body, nil
}

func apiGetJSON(endpoint string, params map[string]string) (map[string]interface{}, error) {
	body, err := apiGet(endpoint, params)
	if err != nil {
		return nil, err
	}
	var result map[string]interface{}
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}
	return result, nil
}

func getToolDefinitions() []map[string]interface{} {
	return []map[string]interface{}{
		{
			"name":        "list_studies",
			"description": "List all imaging studies for this patient.",
			"inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
		},
		{
			"name":        "list_series",
			"description": "List series for a study. Filter by description (AX, T1, FLAIR, etc).",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"study_id": map[string]interface{}{"type": "string", "description": "Study GUID"},
					"filter":   map[string]interface{}{"type": "string", "description": "Filter by description"},
				},
				"required": []string{"study_id"},
			},
		},
		{
			"name":        "list_slices",
			"description": "List slices for a series with position info.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"series_id": map[string]interface{}{"type": "string", "description": "Series GUID"},
				},
				"required": []string{"series_id"},
			},
		},
		{
			"name":        "fetch_image",
			"description": "Fetch slice image as base64 PNG. Optionally set window/level.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"slice_id": map[string]interface{}{"type": "string", "description": "Slice GUID"},
					"wc":       map[string]interface{}{"type": "number", "description": "Window center"},
					"ww":       map[string]interface{}{"type": "number", "description": "Window width"},
				},
				"required": []string{"slice_id"},
			},
		},
		{
			"name":        "query_anatomy",
			"description": "Find slices by region: vertex, pons, medulla, upper_cervical, midbrain, thalamus, ventricles",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"series_id": map[string]interface{}{"type": "string", "description": "Series GUID"},
					"region":    map[string]interface{}{"type": "string", "description": "Anatomical region"},
				},
				"required": []string{"series_id", "region"},
			},
		},
		{
			"name":        "list_lab_tests",
			"description": "List all lab test names for this patient.",
			"inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
		},
		{
			"name":        "get_lab_results",
			"description": "Get lab results. Specify names, optional date range, optional latest flag.",
			"inputSchema": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"names":  map[string]interface{}{"type": "string", "description": "Comma-separated test names"},
					"from":   map[string]interface{}{"type": "string", "description": "Start date YYYY-MM-DD"},
					"to":     map[string]interface{}{"type": "string", "description": "End date YYYY-MM-DD"},
					"latest": map[string]interface{}{"type": "boolean", "description": "Most recent only"},
				},
				"required": []string{"names"},
			},
		},
		{
			"name":        "get_version",
			"description": "Get bridge and server version info.",
			"inputSchema": map[string]interface{}{"type": "object", "properties": map[string]interface{}{}},
		},
	}
}

// Tool handlers
func handleListStudies(id interface{}) {
	body, err := apiGet("/api/studies", nil)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleListSeries(id interface{}, args map[string]interface{}) {
	studyID, _ := args["study_id"].(string)
	if studyID == "" {
		respondError(id, -32602, "study_id required")
		return
	}
	params := map[string]string{"study": studyID}
	if filter, ok := args["filter"].(string); ok && filter != "" {
		params["filter"] = filter
	}
	body, err := apiGet("/api/series", params)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleListSlices(id interface{}, args map[string]interface{}) {
	seriesID, _ := args["series_id"].(string)
	if seriesID == "" {
		respondError(id, -32602, "series_id required")
		return
	}
	body, err := apiGet("/api/slices", map[string]string{"series": seriesID})
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleFetchImage(id interface{}, args map[string]interface{}) {
	sliceID, _ := args["slice_id"].(string)
	if sliceID == "" {
		respondError(id, -32602, "slice_id required")
		return
	}

	// Build image URL with optional W/L params
	u := server + "/image/" + sliceID + "?token=" + token
	if wc, ok := args["wc"].(float64); ok {
		u += fmt.Sprintf("&wc=%.0f", wc)
	}
	if ww, ok := args["ww"].(float64); ok {
		u += fmt.Sprintf("&ww=%.0f", ww)
	}

	log("GET %s", u)
	resp, err := http.Get(u)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		respondError(id, -32000, fmt.Sprintf("HTTP %d: %s", resp.StatusCode, string(body)))
		return
	}

	imgData, err := io.ReadAll(resp.Body)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}

	b64 := base64.StdEncoding.EncodeToString(imgData)
	respondImage(id, b64, "image/png", fmt.Sprintf("Slice %s (%d bytes)", sliceID[:8], len(imgData)))
}

func handleQueryAnatomy(id interface{}, args map[string]interface{}) {
	seriesID, _ := args["series_id"].(string)
	region, _ := args["region"].(string)
	if seriesID == "" || region == "" {
		respondError(id, -32602, "series_id and region required")
		return
	}
	body, err := apiGet("/api/anatomy", map[string]string{"series_id": seriesID, "region": region})
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleListLabTests(id interface{}) {
	body, err := apiGet("/api/labs/tests", nil)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleGetLabResults(id interface{}, args map[string]interface{}) {
	names, _ := args["names"].(string)
	if names == "" {
		respondError(id, -32602, "names required")
		return
	}
	params := map[string]string{"names": names}
	if from, ok := args["from"].(string); ok && from != "" {
		params["from"] = from
	}
	if to, ok := args["to"].(string); ok && to != "" {
		params["to"] = to
	}
	if latest, ok := args["latest"].(bool); ok && latest {
		params["latest"] = "true"
	}
	body, err := apiGet("/api/labs/results", params)
	if err != nil {
		respondError(id, -32000, err.Error())
		return
	}
	respondContent(id, string(body))
}

func handleGetVersion(id interface{}) {
	body, err := apiGet("/api/version", nil)
	if err != nil {
		respondContent(id, fmt.Sprintf("Bridge: %s (built %s)\nServer: unavailable (%v)", Version, BuildTime, err))
		return
	}
	var ver map[string]interface{}
	json.Unmarshal(body, &ver)
	serverVer, _ := ver["version"].(string)
	respondContent(id, fmt.Sprintf("Bridge: %s (built %s)\nServer: %s\nOS: %s/%s", Version, BuildTime, serverVer, runtime.GOOS, runtime.GOARCH))
}

func handleToolCall(id interface{}, params json.RawMessage) {
	var p struct {
		Name      string                 `json:"name"`
		Arguments map[string]interface{} `json:"arguments"`
	}
	json.Unmarshal(params, &p)
	log("Tool: %s args=%v", p.Name, p.Arguments)

	switch p.Name {
	case "list_studies":
		handleListStudies(id)
	case "list_series":
		handleListSeries(id, p.Arguments)
	case "list_slices":
		handleListSlices(id, p.Arguments)
	case "fetch_image":
		handleFetchImage(id, p.Arguments)
	case "query_anatomy":
		handleQueryAnatomy(id, p.Arguments)
	case "list_lab_tests":
		handleListLabTests(id)
	case "get_lab_results":
		handleGetLabResults(id, p.Arguments)
	case "get_version":
		handleGetVersion(id)
	default:
		respondError(id, -32601, "Unknown tool: "+p.Name)
	}
}

func main() {
	log("inou-bridge starting, version=%s, build=%s", Version, BuildTime)

	for _, arg := range os.Args[1:] {
		if strings.HasPrefix(arg, "--server=") {
			server = strings.TrimPrefix(arg, "--server=")
		} else if strings.HasPrefix(arg, "--token=") {
			token = strings.TrimPrefix(arg, "--token=")
		} else if arg == "--version" || arg == "-v" {
			fmt.Printf("inou-bridge %s (built %s)\n", Version, BuildTime)
			os.Exit(0)
		}
	}

	if server == "" || token == "" {
		fmt.Fprintf(os.Stderr, "inou-bridge %s - MCP bridge for Inou medical imaging\n\n", Version)
		fmt.Fprintf(os.Stderr, "Usage: inou-bridge --server=http://host:port --token=PROFILE_GUID\n")
		os.Exit(1)
	}

	log("Server: %s, Token: %s...", server, token[:8])

	scanner := bufio.NewScanner(os.Stdin)
	buf := make([]byte, 64*1024)
	scanner.Buffer(buf, 50*1024*1024)

	for scanner.Scan() {
		line := scanner.Text()
		log("stdin: %d bytes", len(line))

		var req MCPRequest
		if err := json.Unmarshal([]byte(line), &req); err != nil {
			log("Invalid JSON: %v", err)
			continue
		}

		switch req.Method {
		case "initialize":
			respond(req.ID, map[string]interface{}{
				"protocolVersion": "2024-11-05",
				"capabilities":    map[string]interface{}{"tools": map[string]interface{}{}},
				"serverInfo":      map[string]interface{}{"name": "inou-bridge", "version": Version},
			})

		case "notifications/initialized":
			// No response needed

		case "tools/list":
			respond(req.ID, map[string]interface{}{"tools": getToolDefinitions()})

		case "tools/call":
			handleToolCall(req.ID, req.Params)

		default:
			respondError(req.ID, -32601, "Method not found: "+req.Method)
		}
	}

	if err := scanner.Err(); err != nil {
		log("stdin error: %v", err)
	}
	log("stdin closed, exiting")
}