package proxy

import (
	"bufio"
	"crypto/tls"
	"encoding/binary"
	"encoding/hex"
	"math/rand"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	"git.maze.io/maze/styx/internal/log"
)

var seed = rand.NewSource(time.Now().UnixNano())

type Context struct {
	id     int64
	conn   *wrappedConn
	rw     *bufio.ReadWriter
	parent *Session
	data   map[string]any
}

func newContext(conn net.Conn, rw *bufio.ReadWriter, parent *Session) *Context {
	if wrapped, ok := conn.(*wrappedConn); ok {
		conn = wrapped.Conn
	}

	ctx := &Context{
		id:     seed.Int63(),
		conn:   &wrappedConn{Conn: conn},
		rw:     rw,
		parent: parent,
		data:   make(map[string]any),
	}

	return ctx
}

func (ctx *Context) log() log.Logger {
	return log.Console.With().
		Str("context", ctx.ID()).
		Str("addr", ctx.RemoteAddr().String()).
		Logger()
}

func (ctx *Context) ID() string {
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], uint64(ctx.id))
	if ctx.parent != nil {
		return ctx.parent.ID() + "-" + hex.EncodeToString(b[:])
	}
	return hex.EncodeToString(b[:])
}

func (ctx *Context) IsTLS() bool {
	_, ok := ctx.conn.Conn.(*tls.Conn)
	return ok && ctx.parent != nil
}

func (ctx *Context) RemoteAddr() net.Addr {
	if ctx.parent != nil {
		return ctx.parent.ctx.RemoteAddr()
	}
	return ctx.conn.RemoteAddr()
}

func (ctx *Context) SetDeadline(t time.Time) error {
	if ctx.parent != nil {
		return ctx.parent.ctx.SetDeadline(t)
	}
	return ctx.conn.SetDeadline(t)
}

func (ctx *Context) Set(key string, value any) {
	ctx.data[key] = value
}

func (ctx *Context) Get(key string) (value any, ok bool) {
	value, ok = ctx.data[key]
	return
}

func (ctx *Context) Flush() error {
	return ctx.rw.Flush()
}

func (ctx *Context) Write(p []byte) (n int, err error) {
	if n, err = ctx.rw.Write(p); n > 0 {
		atomic.AddInt64(&ctx.conn.bytes, int64(n))
	}
	return
}

type Session struct {
	id       int64
	ctx      *Context
	request  *http.Request
	response *http.Response
	data     map[string]any
}

func newSession(ctx *Context, request *http.Request) *Session {
	return &Session{
		id:      seed.Int63(),
		ctx:     ctx,
		request: request,
		data:    make(map[string]any),
	}
}

func (ses *Session) log() log.Logger {
	return log.Console.With().
		Str("context", ses.ctx.ID()).
		Str("session", ses.ID()).
		Str("addr", ses.ctx.RemoteAddr().String()).
		Logger()
}

func (ses *Session) ID() string {
	var b [8]byte
	binary.BigEndian.PutUint64(b[:], uint64(ses.id))
	return hex.EncodeToString(b[:])
}

func (ses *Session) Context() *Context {
	return ses.ctx
}

func (ses *Session) Request() *http.Request {
	return ses.request
}

func (ses *Session) Response() *http.Response {
	return ses.response
}

type wrappedConn struct {
	net.Conn
	bytes int64
}

func (c *wrappedConn) Write(p []byte) (n int, err error) {
	if n, err = c.Conn.Write(p); n > 0 {
		atomic.AddInt64(&c.bytes, int64(n))
	}
	return
}