package toolsrv

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"syscall"

	"ollie/paths"
)

// RPCResponse is a JSON-RPC 2.0 response.
type RPCResponse struct {
	JSONRPC string          `json:"jsonrpc"`
	ID      json.RawMessage `json:"id"`
	Result  any             `json:"result,omitempty"`
	Error   *RPCError       `json:"error,omitempty"`
	Stream  bool            `json:"stream,omitempty"`
}

// lockedEncoder serializes JSON writes to a single writer.
type lockedEncoder struct {
	mu  sync.Mutex
	enc *json.Encoder
}

func (e *lockedEncoder) Encode(v any) error {
	e.mu.Lock()
	defer e.mu.Unlock()
	return e.enc.Encode(v)
}

// ServeSocket accepts JSON-RPC connections on a Unix socket. Blocks until
// ctx is cancelled.
func (s *Server) ServeSocket(ctx context.Context, sockPath string) error {
	os.Remove(sockPath)
	os.MkdirAll(filepath.Dir(sockPath), 0700)

	ln, err := net.Listen("unix", sockPath)
	if err != nil {
		return fmt.Errorf("listen %s: %w", sockPath, err)
	}
	defer ln.Close()
	defer os.Remove(sockPath)

	fmt.Fprintln(os.Stderr, "ListenReady")

	var mu sync.Mutex
	var activeConns []net.Conn
	var wg sync.WaitGroup

	go func() {
		<-ctx.Done()
		ln.Close()
		mu.Lock()
		for _, c := range activeConns {
			c.Close()
		}
		mu.Unlock()
	}()

	for {
		conn, err := ln.Accept()
		if err != nil {
			if ctx.Err() != nil {
				break
			}
			continue
		}
		mu.Lock()
		activeConns = append(activeConns, conn)
		mu.Unlock()

		wg.Add(1)
		go func() {
			defer wg.Done()
			defer conn.Close()
			s.ServeRPC(ctx, conn, conn)
			mu.Lock()
			for i, c := range activeConns {
				if c == conn {
					activeConns = append(activeConns[:i], activeConns[i+1:]...)
					break
				}
			}
			mu.Unlock()
		}()
	}
	wg.Wait()
	return nil
}

// ServeRPC reads JSON-RPC 2.0 requests from in and writes responses to out.
// Blocks until in is closed or ctx is cancelled.
func (s *Server) ServeRPC(ctx context.Context, in io.Reader, out io.Writer) {
	dec := json.NewDecoder(in)
	enc := &lockedEncoder{enc: json.NewEncoder(out)}

	for {
		var req RPCRequest
		if err := dec.Decode(&req); err != nil {
			if err == io.EOF || ctx.Err() != nil {
				return
			}
			enc.Encode(RPCResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Error:   &RPCError{Code: -32700, Message: "parse error: " + err.Error()},
			})
			continue
		}

		go s.handleRPC(ctx, req, enc)
	}
}

// CWD returns the server's current working directory.
func (s *Server) CWD() string {
	s.wdMu.RLock()
	defer s.wdMu.RUnlock()
	return s.cwd
}

// handleRPC dispatches a single JSON-RPC request.
//
// Reserved method names (control plane):
//   - list_tools, host_info, ping, set_env, set_cwd
//   - detach, list_detached, signal_detached, get_detached_output, dismiss_detached
//   - tool_load
//
// Any other method name is treated as a tool invocation (data plane).
// A tool with the same name as a reserved method is unreachable.
func (s *Server) handleRPC(ctx context.Context, req RPCRequest, enc *lockedEncoder) {
	switch req.Method {
	case "list_tools":
		tools, err := s.ListTools()
		if err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}})
		} else {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: tools})
		}

	case "host_info":
		info := HostInfo{
			Platform:  runtime.GOOS,
			Arch:      runtime.GOARCH,
			IsGitRepo: paths.IsGitRepo(s.CWD()),
		}
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: info})

	case "ping":
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: "pong"})

	case "set_env":
		var params struct {
			Key   string `json:"key"`
			Value string `json:"value"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		s.SetEnv(params.Key, params.Value)
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true})

	case "set_cwd":
		var params struct {
			Dir string `json:"dir"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		s.SetCWD(params.Dir)
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true})

	case "detach":
		ok := s.Detach()
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok})

	case "list_detached":
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: s.ListDetachedRaw()})

	case "signal_detached":
		var params struct {
			PID    int `json:"pid"`
			Signal int `json:"signal"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		if err := s.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}})
		} else {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true})
		}

	case "tool_load":
		var params struct {
			Name string `json:"name"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		if err := s.LoadTool(params.Name); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}})
		} else {
			res, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name})
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(res)})
		}

	case "get_detached_output":
		var params struct {
			PID int `json:"pid"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		out, err := s.GetDetachedOutput(params.PID)
		if err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}})
		} else {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: out})
		}

	case "dismiss_detached":
		var params struct {
			PID int `json:"pid"`
		}
		if err := json.Unmarshal(req.Params, &params); err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}})
			return
		}
		ok := s.DismissDetached(params.PID)
		enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok})

	default:
		// Data plane: method name is the tool name.
		streamCtx := WithOutputStream(ctx, func(data string) {
			enc.Encode(RPCResponse{
				JSONRPC: "2.0",
				ID:      req.ID,
				Result:  OutputNotification{Data: data},
				Stream:  true,
			})
		})
		result, err := s.CallTool(streamCtx, req.Method, req.Params)
		if err != nil {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}})
		} else {
			enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(result)})
		}
	}
}