package main

import (
 "bytes"
 "encoding/json"
 "fmt"
 "io"
 "log"
 "net/http"
 "os"
 "path/filepath"
 "strconv"
 "strings"
 "time"
)

// --- MCP Protocol Types ---
type MCPRequest struct {
 JSONRPC string          `json:"jsonrpc"`
 ID      any             `json:"id,omitempty"`
 Method  string          `json:"method"`
 Params  json.RawMessage `json:"params,omitempty"`
}

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

type MCPContent struct {
 Type string `json:"type"`
 Text string `json:"text"`
}

func sendMCPResponse(id any, result any) MCPResponse {
 return MCPResponse{JSONRPC: "2.0", ID: id, Result: result}
}

func (s *TaoServer) dispatchMCP(token string, client string, req MCPRequest) {
 var resp MCPResponse
 resp.JSONRPC = "2.0"
 resp.ID = req.ID

 switch req.Method {
 case "initialize":
  resp.Result = map[string]interface{}{
   "protocolVersion": "2025-06-18",
   "capabilities": map[string]interface{}{
    "tools":     map[string]interface{}{"listChanged": false},
    "resources": map[string]interface{}{"listChanged": false},
    "prompts":   map[string]interface{}{"listChanged": false},
    "logging":   map[string]interface{}{},
   },
   "serverInfo": map[string]string{
    "name":    "Tao-Memory-Server",
    "version": "1.2.0",
   },
  }
 case "notifications/initialized":
  log.Printf("[MCP Notify] initialized from %s", token)
  return
 case "tools/list":
  resp.Result = map[string]interface{}{
   "tools": buildToolList(),
  }
 case "tools/call":
  var params struct {
   Name      string                 `json:"name"`
   Arguments map[string]interface{} `json:"arguments"`
  }
  _ = json.Unmarshal(req.Params, &params)
  if tool, ok := ToolRegistry[params.Name]; ok {
   result, err := tool.Handler(params.Arguments)
   if err != nil {
    resp.Result = map[string]interface{}{
     "content": []MCPContent{{Type: "text", Text: "error: " + err.Error()}},
    }
   } else {
    resp.Result = map[string]interface{}{
     "content": []MCPContent{{Type: "text", Text: result}},
    }
   }
  } else {
   resp.Result = map[string]interface{}{
    "content": []MCPContent{{Type: "text", Text: "error: tool not found"}},
   }
  }
 default:
  _ = s.Record("agent_action", fmt.Sprintf("执行指令: %+v", req), 2)
  return
 }

 if token == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
  log.Printf("[MCP Response] missing token for method=%s", req.Method)
  return
 }
 connKey := buildConnKey(token, client)
 if ch, ok := s.conns.Load(connKey); ok {
  if b, err := json.Marshal(resp); err == nil {
   ch.(chan string) <- string(b)
   log.Printf("[MCP Response] sent via SSE method=%s", req.Method)
  }
 } else {
  log.Printf("[MCP Response] no SSE channel for token=%s client=%s method=%s", token, client, req.Method)
 }
}

func getEnv(key, def string) string {
 if v := os.Getenv(key); v != "" {
  return v
 }
 return def
}

func getEnvBool(key string, def bool) bool {
 v := strings.ToLower(strings.TrimSpace(os.Getenv(key)))
 if v == "" {
  return def
 }
 switch v {
 case "1", "true", "yes", "on":
  return true
 case "0", "false", "no", "off":
  return false
 default:
  return def
 }
}

func getEnvInt(key string, def int) int {
 if v := os.Getenv(key); v != "" {
  if n, err := strconv.Atoi(v); err == nil {
   return n
  }
 }
 return def
}

func extractToken(r *http.Request) (string, bool) {
 if q := r.URL.Query().Get("token"); q != "" {
  return q, true
 }
 h := r.Header.Get("Authorization")
 if strings.HasPrefix(h, "Bearer ") {
  return strings.TrimSpace(strings.TrimPrefix(h, "Bearer ")), false
 }
 return "", false
}

func buildConnKey(token string, client string) string {
 if token == "" {
  token = "anon"
 }
 if client != "" {
  return token + "_" + client
 }
 return token
}

func generateClientID() string {
 return fmt.Sprintf("c%d", time.Now().UnixNano())
}

func parseCORSOrigins() (bool, []string) {
 raw := strings.TrimSpace(os.Getenv("TAO_CORS_ORIGINS"))
 if raw == "" {
  return false, nil
 }
 if raw == "*" {
  return true, nil
 }
 parts := strings.Split(raw, ",")
 var origins []string
 for _, p := range parts {
  p = strings.TrimSpace(p)
  if p != "" {
   origins = append(origins, p)
  }
 }
 return false, origins
}

func setCORSHeaders(w http.ResponseWriter, r *http.Request) {
 origin := r.Header.Get("Origin")
 if origin == "" {
  return
 }
 allowAll, origins := parseCORSOrigins()
 if allowAll {
  w.Header().Set("Access-Control-Allow-Origin", "*")
 } else {
  allowed := false
  for _, o := range origins {
   if o == origin {
    allowed = true
    break
   }
  }
  if !allowed {
   return
  }
  w.Header().Set("Access-Control-Allow-Origin", origin)
  w.Header().Set("Vary", "Origin")
 }
 w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
 w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
}

func isSubpath(path string, base string) bool {
 absPath, err1 := filepath.Abs(path)
 absBase, err2 := filepath.Abs(base)
 if err1 != nil || err2 != nil {
  return false
 }
 if absPath == absBase {
  return true
 }
 return strings.HasPrefix(absPath, absBase+string(filepath.Separator))
}

// --- 以简御繁：鉴权 ---
func (s *TaoServer) checkAuth(r *http.Request) bool {
 token := getEnv("TAO_AUTH_TOKEN", "")
 if token == "" {
  return getEnvBool("TAO_ALLOW_ANON", false)
 }
 reqToken, _ := extractToken(r)
 if reqToken == token {
  return true
 }
 return false
}

func (s *TaoServer) requireAuth(next http.HandlerFunc) http.HandlerFunc {
 return func(w http.ResponseWriter, r *http.Request) {
  if r.Method == "OPTIONS" {
   setCORSHeaders(w, r)
   w.WriteHeader(http.StatusOK)
   return
  }
  if !s.checkAuth(r) {
   http.Error(w, "Unauthorized", http.StatusUnauthorized)
   return
  }
  next(w, r)
 }
}

// --- 感官 (Webhook Adapters) ---
// 适配 Gitea 的 Push Webhook
func (s *TaoServer) GiteaHandler(w http.ResponseWriter, r *http.Request) {
 var payload struct {
  Repository struct {
   Name string `json:"name"`
  } `json:"repository"`
  Commits []struct {
   Message string `json:"message"`
  } `json:"commits"`
 }

 if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
  http.Error(w, "Bad Request", 400)
  return
 }

 if len(payload.Commits) > 0 {
  msg := payload.Commits[0].Message
  summary := fmt.Sprintf("代码演化于 [%s]: %s", payload.Repository.Name, msg)
  _ = s.Record("code", summary, 4)
 }
 w.WriteHeader(200)
}

// 适配 SmsReceiver-go 的短信推送
func (s *TaoServer) SmsHandler(w http.ResponseWriter, r *http.Request) {
 var payload struct {
  From    string `json:"from"`
  Content string `json:"content"`
 }
 if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
  http.Error(w, "Bad Request", 400)
  return
 }

 summary := fmt.Sprintf("收到信号 [%s]: %s", payload.From, payload.Content)
 _ = s.Record("sms", summary, 3)
 w.WriteHeader(200)
}

// --- MCP SSE ---
func (s *TaoServer) SSEHandler(w http.ResponseWriter, r *http.Request) {
 log.Printf("[SSE Connect] Remote=%s URL=%s", r.RemoteAddr, r.URL.String())
 w.Header().Set("Content-Type", "text/event-stream")
 w.Header().Set("Cache-Control", "no-cache")
 w.Header().Set("Connection", "keep-alive")
 w.Header().Set("X-Accel-Buffering", "no")
 setCORSHeaders(w, r)

 flusher, ok := w.(http.Flusher)
 if !ok {
  http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
  return
 }

 // 告知客户端 POST 入口（按客户端拼接习惯输出）
 style := getEnv("TAO_ENDPOINT_STYLE", "message")
 endpoint := "mcp/message"
 if style == "message" {
  endpoint = "message"
 }

 queryToken := r.URL.Query().Get("token")
 token, _ := extractToken(r)
 if token == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
  http.Error(w, "Unauthorized", http.StatusUnauthorized)
  return
 }

 client := r.URL.Query().Get("client")
 if client == "" {
  client = generateClientID()
 }

 endpointToken := queryToken
 if endpointToken == "" && getEnvBool("TAO_ENDPOINT_TOKEN_ECHO", false) {
  endpointToken = token
 }

 if endpointToken != "" {
  if strings.Contains(endpoint, "?") {
   endpoint = endpoint + "&token=" + endpointToken
  } else {
   endpoint = endpoint + "?token=" + endpointToken
  }
 }
 if client != "" {
  if strings.Contains(endpoint, "?") {
   endpoint = endpoint + "&client=" + client
  } else {
   endpoint = endpoint + "?client=" + client
  }
 }
 fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)
 flusher.Flush()

 msgChan := make(chan string, 50)
 connKey := buildConnKey(token, client)
 s.conns.Store(connKey, msgChan)
 defer s.conns.Delete(connKey)

 ticker := time.NewTicker(5 * time.Second)
 defer ticker.Stop()

 for {
  select {
  case <-r.Context().Done():
   return
  case <-ticker.C:
   fmt.Fprintf(w, ":ping\n\n")
   flusher.Flush()
  case msg := <-msgChan:
   fmt.Fprintf(w, "event: message\ndata: %s\n\n", msg)
   flusher.Flush()
  }
 }
}

// --- MCP Message ---
func (s *TaoServer) MessageHandler(w http.ResponseWriter, r *http.Request) {
 setCORSHeaders(w, r)

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

 bodyBytes, _ := io.ReadAll(r.Body)
 r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
 if getEnvBool("TAO_DEBUG", false) {
  log.Printf("[MCP POST] From=%s URL=%s Body=%s", r.RemoteAddr, r.URL.String(), string(bodyBytes))
 } else {
  log.Printf("[MCP POST] From=%s URL=%s", r.RemoteAddr, r.URL.String())
 }

 var req MCPRequest
 if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
  http.Error(w, "Bad Request", 400)
  return
 }

 token, _ := extractToken(r)
 client := r.URL.Query().Get("client")
 w.WriteHeader(http.StatusAccepted)

 go s.dispatchMCP(token, client, req)
}

// --- MCP Unified Gateway (/mcp) ---
func (s *TaoServer) MCPUnifiedHandler(w http.ResponseWriter, r *http.Request) {
 if r.Method == "OPTIONS" {
  w.Header().Set("Access-Control-Allow-Origin", "*")
  w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
  w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
  w.WriteHeader(http.StatusOK)
  return
 }
 if r.Method == http.MethodGet {
  s.SSEHandler(w, r)
  return
 }
 if r.Method == http.MethodPost {
  accept := r.Header.Get("Accept")
  if strings.Contains(accept, "text/event-stream") {
   s.SSEHandler(w, r)
   return
  }
  if r.ContentLength == 0 {
   s.SSEHandler(w, r)
   return
  }
  bodyBytes, _ := io.ReadAll(r.Body)
  r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
  if len(bodyBytes) == 0 {
   s.SSEHandler(w, r)
   return
  }
  if strings.Contains(string(bodyBytes), "\"jsonrpc\":\"2.0\"") {
   s.MessageHandler(w, r)
   return
  }
  http.Error(w, "Bad Request", 400)
  return
 }
 http.Error(w, "Method Not Allowed", 405)
}

// --- 主程序 (Main) ---
func main() {
 if getEnv("TAO_AUTH_TOKEN", "") == "" && !getEnvBool("TAO_ALLOW_ANON", false) {
  log.Fatal("TAO_AUTH_TOKEN is required unless TAO_ALLOW_ANON=true")
 }

 memoryRoot := getEnv("MEMORY_ROOT", "./knowledge_ocean")
 searchRoot := getEnv("TAO_SEARCH_ROOT", memoryRoot)
 if !isSubpath(searchRoot, memoryRoot) {
  log.Printf("TAO_SEARCH_ROOT must be under MEMORY_ROOT, fallback to MEMORY_ROOT")
  searchRoot = memoryRoot
 }

 server := &TaoServer{
  config: Config{
   MemoryRoot:    memoryRoot,
   Port:          getEnv("PORT", "5001"),
   SearchRoot:    searchRoot,
   MaxSearchFiles: getEnvInt("TAO_SEARCH_MAX_FILES", 2000),
  },
 }

 server.RegisterTools()

 // 启动 Webhook 监听 (感知层)
 http.HandleFunc("/ingest/gitea", server.requireAuth(server.GiteaHandler))
 http.HandleFunc("/ingest/sms", server.requireAuth(server.SmsHandler))

 // MCP SSE + Message
 http.HandleFunc("/mcp/sse", server.requireAuth(server.SSEHandler))
 http.HandleFunc("/mcp", server.requireAuth(server.MCPUnifiedHandler))
 http.HandleFunc("/mcp/", server.requireAuth(server.MCPUnifiedHandler))
 http.HandleFunc("/mcp/message", server.requireAuth(server.MessageHandler))

 fmt.Printf("Tao Memory Server 启动。道场地址: :%s\n", server.config.Port)
 log.Fatal(http.ListenAndServe(":"+server.config.Port, nil))
}