package serial

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"time"
)

const (
	streamProtocolVersion      = 1
	maxStreamHandshakeBytes    = 4 * 1024
	defaultStreamSocketTimeout = 10 * time.Second
)

type streamLeaseRequest struct {
	Version  int    `json:"version"`
	Action   string `json:"action"`
	Channel  string `json:"channel"`
	ClientID string `json:"client_id"`
	LeaseID  string `json:"lease_id,omitempty"`
}

type streamLeaseResponse struct {
	Version int             `json:"version"`
	OK      bool            `json:"ok"`
	Channel string          `json:"channel,omitempty"`
	LeaseID string          `json:"lease_id,omitempty"`
	Stats   *StreamSnapshot `json:"stats,omitempty"`
	Code    string          `json:"code,omitempty"`
	Error   string          `json:"error,omitempty"`
}

// StreamSocketServer exposes a separate data-plane socket. Unix socket mode
// and service-group ownership authenticate tera-sessiond; this protocol never
// accepts a raw /dev path.
type StreamSocketServer struct {
	Manager          *StreamLeaseManager
	HandshakeTimeout time.Duration
}

func (s StreamSocketServer) Serve(listener net.Listener) error {
	if s.Manager == nil {
		return errors.New("serial stream lease manager is required")
	}
	for {
		connection, err := listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return nil
			}
			if temporary, ok := err.(interface{ Temporary() bool }); ok && temporary.Temporary() {
				continue
			}
			return err
		}
		go s.handleConnection(connection)
	}
}

func (s StreamSocketServer) handleConnection(connection net.Conn) {
	defer connection.Close()
	timeout := s.HandshakeTimeout
	if timeout <= 0 {
		timeout = defaultStreamSocketTimeout
	}
	_ = connection.SetDeadline(time.Now().Add(timeout))
	reader := bufio.NewReader(io.LimitReader(connection, maxStreamHandshakeBytes+1))
	payload, err := reader.ReadBytes('\n')
	if err != nil || len(payload) > maxStreamHandshakeBytes || !bytes.HasSuffix(payload, []byte{'\n'}) {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "invalid serial stream handshake"})
		return
	}
	if reader.Buffered() != 0 {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial bytes cannot precede lease acceptance"})
		return
	}
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	var request streamLeaseRequest
	if err := decoder.Decode(&request); err != nil {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "invalid serial stream handshake"})
		return
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial stream handshake must contain one JSON value"})
		return
	}
	if request.Version != streamProtocolVersion {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial stream protocol is not supported"})
		return
	}
	if request.Action == "status" {
		if request.Channel != "" || request.ClientID != "" || request.LeaseID != "" {
			s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial stream status must not include a lease target"})
			return
		}
		status := s.Manager.Snapshot()
		s.writeResponse(connection, streamLeaseResponse{OK: true, Stats: &status})
		return
	}
	if request.Action == "break" {
		if request.Channel == "" || request.ClientID == "" || request.LeaseID == "" {
			s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial stream break requires an active lease"})
			return
		}
		if err := s.Manager.SendBreak(context.Background(), request.Channel, request.LeaseID, request.ClientID); err != nil {
			s.writeResponse(connection, streamLeaseResponse{Code: "hardware_failed", Error: err.Error()})
			return
		}
		s.writeResponse(connection, streamLeaseResponse{OK: true, Channel: request.Channel, LeaseID: request.LeaseID})
		return
	}
	if request.Action != "lease" {
		s.writeResponse(connection, streamLeaseResponse{Code: "validation_failed", Error: "serial stream action is not supported"})
		return
	}
	lease, err := s.Manager.Acquire(context.Background(), request.Channel, request.ClientID)
	if err != nil {
		s.writeLeaseError(connection, err)
		return
	}
	defer lease.Close()
	if !s.writeResponse(connection, streamLeaseResponse{
		Version: streamProtocolVersion, OK: true,
		Channel: request.Channel, LeaseID: lease.LeaseID(),
	}) {
		return
	}
	_ = connection.SetDeadline(time.Time{})

	done := make(chan struct{}, 2)
	go func() {
		_, _ = io.Copy(lease, connection)
		done <- struct{}{}
	}()
	go func() {
		_, _ = io.Copy(connection, lease)
		done <- struct{}{}
	}()
	<-done
	_ = lease.Close()
	_ = connection.Close()
	<-done
}

func (s StreamSocketServer) writeLeaseError(connection net.Conn, err error) {
	response := streamLeaseResponse{Version: streamProtocolVersion, Code: "hardware_failed", Error: err.Error()}
	switch {
	case errors.Is(err, ErrStreamLeaseBusy):
		response.Code = "lease_busy"
	case IsUnavailable(err):
		response.Code = "hardware_unavailable"
	default:
		response.Code = "validation_failed"
	}
	s.writeResponse(connection, response)
}

func (StreamSocketServer) writeResponse(connection net.Conn, response streamLeaseResponse) bool {
	if response.Version == 0 {
		response.Version = streamProtocolVersion
	}
	return json.NewEncoder(connection).Encode(response) == nil
}

// StreamSocketClient implements directstream.TransportFactory without
// exposing a device path outside tera-seriald.
type StreamSocketClient struct {
	path     string
	clientID string
	timeout  time.Duration
}

func NewStreamSocketClient(path, clientID string) *StreamSocketClient {
	if path == "" {
		path = DefaultStreamSocketPath
	}
	if clientID == "" {
		clientID = "tera-sessiond"
	}
	return &StreamSocketClient{path: path, clientID: clientID, timeout: defaultStreamSocketTimeout}
}

func (c *StreamSocketClient) Open(ctx context.Context, channel string) (io.ReadWriteCloser, error) {
	dialer := net.Dialer{Timeout: c.timeout}
	connection, err := dialer.DialContext(ctx, "unix", c.path)
	if err != nil {
		return nil, &UnavailableError{Reason: "tera-seriald stream socket is not reachable"}
	}
	fail := true
	defer func() {
		if fail {
			_ = connection.Close()
		}
	}()
	deadline := time.Now().Add(c.timeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	_ = connection.SetDeadline(deadline)
	request := streamLeaseRequest{
		Version: streamProtocolVersion, Action: "lease", Channel: channel, ClientID: c.clientID,
	}
	if err := json.NewEncoder(connection).Encode(request); err != nil {
		return nil, errors.New("serial stream handshake failed")
	}
	reader := bufio.NewReader(connection)
	payload, err := reader.ReadBytes('\n')
	if err != nil || len(payload) > maxStreamHandshakeBytes {
		return nil, &UnavailableError{Reason: "tera-seriald returned no stream lease"}
	}
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	var response streamLeaseResponse
	if err := decoder.Decode(&response); err != nil {
		return nil, &UnavailableError{Reason: "tera-seriald returned an invalid stream lease"}
	}
	if response.Version != streamProtocolVersion {
		return nil, errors.New("tera-seriald stream protocol version mismatch")
	}
	if !response.OK {
		switch response.Code {
		case "lease_busy":
			return nil, fmt.Errorf("%w: %s", ErrStreamLeaseBusy, response.Error)
		case "hardware_unavailable":
			return nil, &UnavailableError{Reason: response.Error}
		default:
			return nil, errors.New(response.Error)
		}
	}
	if response.Channel != channel || response.LeaseID == "" {
		return nil, errors.New("tera-seriald returned a mismatched stream lease")
	}
	_ = connection.SetDeadline(time.Time{})
	fail = false
	return &streamSocketTransport{
		Conn: connection, reader: reader, path: c.path, channel: channel,
		clientID: c.clientID, leaseID: response.LeaseID, timeout: c.timeout,
	}, nil
}

func (c *StreamSocketClient) Status(ctx context.Context) (StreamSnapshot, error) {
	dialer := net.Dialer{Timeout: c.timeout}
	connection, err := dialer.DialContext(ctx, "unix", c.path)
	if err != nil {
		return StreamSnapshot{}, &UnavailableError{Reason: "tera-seriald stream socket is not reachable"}
	}
	defer connection.Close()
	deadline := time.Now().Add(c.timeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	_ = connection.SetDeadline(deadline)
	if err := json.NewEncoder(connection).Encode(streamLeaseRequest{Version: streamProtocolVersion, Action: "status"}); err != nil {
		return StreamSnapshot{}, errors.New("serial stream status request failed")
	}
	var response streamLeaseResponse
	if err := json.NewDecoder(io.LimitReader(connection, maxStreamHandshakeBytes)).Decode(&response); err != nil {
		return StreamSnapshot{}, &UnavailableError{Reason: "tera-seriald returned no stream status"}
	}
	if response.Version != streamProtocolVersion || !response.OK || response.Stats == nil {
		return StreamSnapshot{}, errors.New("tera-seriald returned an invalid stream status")
	}
	return *response.Stats, nil
}

type streamSocketTransport struct {
	net.Conn
	reader   *bufio.Reader
	path     string
	channel  string
	clientID string
	leaseID  string
	timeout  time.Duration
}

func (t *streamSocketTransport) Read(buffer []byte) (int, error) {
	return t.reader.Read(buffer)
}

func (t *streamSocketTransport) SendBreak(ctx context.Context) error {
	dialer := net.Dialer{Timeout: t.timeout}
	connection, err := dialer.DialContext(ctx, "unix", t.path)
	if err != nil {
		return &UnavailableError{Reason: "tera-seriald stream socket is not reachable"}
	}
	defer connection.Close()
	deadline := time.Now().Add(t.timeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	_ = connection.SetDeadline(deadline)
	request := streamLeaseRequest{Version: streamProtocolVersion, Action: "break", Channel: t.channel, ClientID: t.clientID, LeaseID: t.leaseID}
	if err := json.NewEncoder(connection).Encode(request); err != nil {
		return errors.New("serial stream break request failed")
	}
	var response streamLeaseResponse
	if err := json.NewDecoder(io.LimitReader(connection, maxStreamHandshakeBytes)).Decode(&response); err != nil {
		return &UnavailableError{Reason: "tera-seriald returned no break response"}
	}
	if response.Version != streamProtocolVersion || !response.OK || response.Channel != t.channel || response.LeaseID != t.leaseID {
		if response.Error != "" {
			return errors.New(response.Error)
		}
		return errors.New("tera-seriald returned an invalid break response")
	}
	return nil
}