package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"net/http/httputil"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/mcp"
	"g.fc.run/ai/byre/internal/store"
)

const maxPreviewName = 40

type exposePortArgs struct {
	Name string `json:"name"`
	Port int    `json:"port"`
}

func (m *Manager) exposePortTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "expose_port",
		Description: "Show the human a server you are running, such as a dev server, as a link above the task's tabs that opens it in the task's Page view beside its transcript. " +
			"Start the server yourself first (in the background), listening on 127.0.0.1 (not 0.0.0.0 or localhost) on any free port. " +
			"byre then serves it on the same port of byre's tailnet node, and the human's app opens it through a local link. " +
			"name labels it in the tab (\"web\", \"storybook\"); exposing a name or port again replaces the earlier one. " +
			"The server stops when your process does (a byre restart, compaction, the task closing); the human sees it as unreachable until you start it again.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"name": {"type": "string", "description": "Short label, shown as the link's text."},
				"port": {"type": "integer", "description": "The port the server listens on, at 127.0.0.1."}
			},
			"required": ["name", "port"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.exposePort(ctx, l, raw)
		},
	}
}

func (m *Manager) exposePort(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := exposePortArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed expose_port call: %w", err)
	}

	name := strings.TrimSpace(args.Name)
	if name == "" || strings.ContainsAny(name, "\r\n") || utf8.RuneCountInString(name) > maxPreviewName {
		return "", fmt.Errorf("name must be one non-empty line of at most %d characters", maxPreviewName)
	}

	if args.Port < 1 || args.Port > 65535 {
		return "", fmt.Errorf("port %d is not a TCP port", args.Port)
	}

	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(args.Port))
	conn, err := net.DialTimeout("tcp", addr, time.Second)
	if err != nil {
		return "", fmt.Errorf("nothing is listening on %s; start the server listening on 127.0.0.1 first: %w", addr, err)
	}
	conn.Close()

	m.exposeMu.Lock()
	defer m.exposeMu.Unlock()
	m.previewMu.Lock()
	serving := m.previewListeners[args.Port] != nil
	m.previewMu.Unlock()
	if !serving {
		err = m.servePreview(args.Port)
		if err != nil {
			return "", err
		}
	}

	dropped, err := m.store.SetPreview(ctx, api.Preview{TaskId: l.taskID, Name: name, Port: args.Port, CreatedAt: time.Now().UTC()}, agentActor)
	if errors.Is(err, store.ErrNotFound) {
		if !serving {
			m.stopPreview(args.Port)
		}
		return "", errors.New("this task is no longer active")
	}
	if err != nil {
		return "", err
	}
	for _, port := range dropped {
		m.stopPreview(port)
	}

	return fmt.Sprintf("Exposed %q: port %d of byre's tailnet node proxies to %s, and the human sees a link above the task's tabs that opens it in the Page view beside its transcript.", name, args.Port, addr), nil
}

func (m *Manager) servePreview(port int) error {
	m.stopPreview(port)
	ln, err := m.tailnet.Listen("tcp", ":"+strconv.Itoa(port))
	if err != nil {
		return fmt.Errorf("listen on the tailnet on port %d: %w", port, err)
	}
	lc, err := m.tailnet.LocalClient()
	if err != nil {
		ln.Close()
		return err
	}
	// HTTP previews also arrive through the app's loopback forwards. Reject
	// attacker-controlled hostnames before proxying: DNS rebinding makes those
	// requests same-origin, so cross-origin protection alone cannot stop them.
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	status, err := lc.StatusWithoutPeers(ctx)
	if err != nil {
		ln.Close()
		return fmt.Errorf("get preview hostnames: %w", err)
	}
	hosts := map[string]bool{}
	if status.Self != nil {
		name := strings.ToLower(strings.TrimSuffix(status.Self.DNSName, "."))
		hosts[name] = true
		hosts[strings.SplitN(name, ".", 2)[0]] = true
	}
	for _, ip := range status.TailscaleIPs {
		hosts[ip.String()] = true
	}

	tracked := &previewListener{Listener: ln, conns: map[net.Conn]bool{}}
	m.previewMu.Lock()
	m.previewListeners[port] = tracked
	m.previewMu.Unlock()

	target := net.JoinHostPort("127.0.0.1", strconv.Itoa(port))
	proxy := &httputil.ReverseProxy{
		Rewrite: func(pr *httputil.ProxyRequest) {
			pr.Out.URL.Scheme = "http"
			pr.Out.URL.Host = target
		},
		FlushInterval: -1,
	}
	protection := http.NewCrossOriginProtection()
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		check := r
		if r.Header.Get("Upgrade") != "" {
			// An upgraded GET opens a bidirectional channel, so it needs the
			// same origin checks as a state-changing request.
			upgrade := *r
			upgrade.Method = http.MethodPost
			check = &upgrade
		}
		if err := protection.Check(check); err != nil {
			http.Error(w, err.Error(), http.StatusForbidden)
			return
		}
		host := r.Host
		if h, _, err := net.SplitHostPort(host); err == nil {
			host = h
		}
		host = strings.ToLower(strings.TrimSuffix(host, "."))
		ip := net.ParseIP(host)
		if host == "" || !(hosts[host] || host == "localhost" || strings.HasSuffix(host, ".localhost") || (ip != nil && ip.IsLoopback())) {
			http.Error(w, "forbidden preview host", http.StatusForbidden)
			return
		}
		who, err := lc.WhoIs(r.Context(), r.RemoteAddr)
		if err != nil {
			http.Error(w, err.Error(), http.StatusForbidden)
			return
		}
		if who.Node.IsTagged() {
			http.Error(w, "tagged nodes are not users", http.StatusForbidden)
			return
		}
		r.Header.Set("Tailscale-User-Login", who.UserProfile.LoginName)
		r.Header.Set("Tailscale-User-Name", who.UserProfile.DisplayName)
		proxy.ServeHTTP(w, r)
	})
	go func() {
		err := (&http.Server{Handler: handler}).Serve(tracked)
		if !errors.Is(err, net.ErrClosed) {
			slog.Error("serve preview", "port", port, "err", err)
		}
	}()
	return nil
}

func (m *Manager) stopPreview(port int) {
	m.previewMu.Lock()
	ln := m.previewListeners[port]
	delete(m.previewListeners, port)
	m.previewMu.Unlock()
	if ln != nil {
		ln.closeAll()
	}
}

type previewListener struct {
	net.Listener
	mu     sync.Mutex
	conns  map[net.Conn]bool
	closed bool
}

func (l *previewListener) Accept() (net.Conn, error) {
	c, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.closed {
		c.Close()
		return nil, net.ErrClosed
	}
	tc := &previewConn{Conn: c, l: l}
	l.conns[tc] = true
	return tc, nil
}

func (l *previewListener) closeAll() {
	l.Listener.Close()
	l.mu.Lock()
	l.closed = true
	conns := l.conns
	l.conns = map[net.Conn]bool{}
	l.mu.Unlock()
	for c := range conns {
		c.Close()
	}
}

type previewConn struct {
	net.Conn
	l *previewListener
}

func (c *previewConn) Close() error {
	c.l.mu.Lock()
	delete(c.l.conns, c)
	c.l.mu.Unlock()
	return c.Conn.Close()
}

func (m *Manager) recoverPreviews(ctx context.Context) error {
	previews, err := m.store.ActivePreviews(ctx)
	if err != nil {
		return err
	}
	for _, p := range previews {
		err = m.servePreview(p.Port)
		if err != nil {
			slog.Error("serve preview", "task", p.TaskId, "port", p.Port, "err", err)
		}
	}
	return nil
}

func (m *Manager) Unexpose(ctx context.Context, taskID string) {
	m.exposeMu.Lock()
	defer m.exposeMu.Unlock()
	previews, err := m.store.ListPreviews(ctx, taskID)
	if err != nil {
		slog.Error("list previews", "task", taskID, "err", err)
		return
	}
	active, err := m.store.ActivePreviews(ctx)
	if err != nil {
		slog.Error("list active previews", "err", err)
		return
	}
	held := map[int]bool{}
	for _, p := range active {
		if p.TaskId != taskID {
			held[p.Port] = true
		}
	}
	for _, p := range previews {
		if held[p.Port] {
			continue
		}
		m.stopPreview(p.Port)
		slog.Info("stopped serving preview", "task", taskID, "name", p.Name, "port", p.Port)
	}
}