package tcp

import (
	"encoding/json"
	"io"
	"net"

	"go.uber.org/zap"
)

type Handler func(w ResponseWriter, req *Request) error
type Middleware func(next Handler) Handler

type Router struct {
	middlewares []Middleware
	routes      map[Method]map[string]Handler
}

func NewRouter() *Router {
	return &Router{
		routes: make(map[Method]map[string]Handler),
	}
}

func (self *Router) Use(middleware Middleware) *Router {
	self.middlewares = append(self.middlewares, middleware)
	return self
}

func (self *Router) Register(method Method, route string, handler Handler) {
	_, ok := self.routes[method]
	if !ok {
		self.routes[method] = make(map[string]Handler)
	}

	for _, middleware := range self.middlewares {
		handler = middleware(handler)
	}
	self.routes[method][route] = handler
}

func (self *Router) GET(route string, handler Handler) {
	self.Register(MethodGET, route, handler)
}

func (self *Router) POST(route string, handler Handler) {
	self.Register(MethodPOST, route, handler)
}

func (self *Router) PUT(route string, handler Handler) {
	self.Register(MethodPUT, route, handler)
}

func (self *Router) DELETE(route string, handler Handler) {
	self.Register(MethodDELETE, route, handler)
}

func (self *Router) SOCKET(route string, handler Handler) {
	self.Register(MethodSOCKET, route, handler)
}

func (self *Router) run(conn net.Conn, req *Request) {
	handler, ok := self.routes[req.Method][req.Route]
	if !ok {
		zap.L().Warn("route not exist",
			zap.String("route", req.Route))
		return
	}

	w := NewResponseBuilder(conn)
	err := handler(w, req)
	if err != nil {
		zap.L().Error("failed to run handler",
			zap.Error(err))
		return
	}

	res := w.Build(req)
	header, err := res.Header()
	if err != nil {
		zap.L().Error("failed to marshal header",
			zap.Error(err))
		return
	}
	if err := sendFrame(conn, header); err != nil {
		zap.L().Error("failed to write header",
			zap.Error(err))
		return
	}
	if err := sendFrame(conn, res.Body); err != nil {
		zap.L().Error("failed to write body",
			zap.Error(err))
		return
	}
}

func (self *Router) Listen(addr string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			return err
		}

		rawHeader, err := readFrame(conn)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}

		body, err := readFrame(conn)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}

		var header RequestHeader
		if err := json.Unmarshal(rawHeader, &header); err != nil {
			return err
		}
		req := NewRequest(conn, header, body)
		go self.run(conn, req)
	}

	return nil
}