package runner

import (
	"fmt"
	"sync"

	"github.com/valyala/fasthttp"

	"Moonshark/core/utils/logger"

	"maps"

	luajit "git.sharkk.net/Sky/LuaJIT-to-Go"
)

// Error represents a simple error string
type Error string

func (e Error) Error() string {
	return string(e)
}

// Error types
var (
	ErrSandboxNotInitialized = Error("sandbox not initialized")
)

// Sandbox provides a secure execution environment for Lua scripts
type Sandbox struct {
	modules map[string]any
	debug   bool
	mu      sync.RWMutex
}

// NewSandbox creates a new sandbox environment
func NewSandbox() *Sandbox {
	return &Sandbox{
		modules: make(map[string]any, 8),
		debug:   false,
	}
}

// EnableDebug turns on debug logging
func (s *Sandbox) EnableDebug() {
	s.debug = true
}

// debugLog logs a message if debug mode is enabled
func (s *Sandbox) debugLog(format string, args ...interface{}) {
	if s.debug {
		logger.Debug("Sandbox "+format, args...)
	}
}

// AddModule adds a module to the sandbox environment
func (s *Sandbox) AddModule(name string, module any) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.modules[name] = module
	s.debugLog("Added module: %s", name)
}

// Setup initializes the sandbox in a Lua state
func (s *Sandbox) Setup(state *luajit.State) error {
	s.debugLog("Setting up sandbox...")

	// Load the sandbox code
	if err := loadSandboxIntoState(state); err != nil {
		s.debugLog("Failed to load sandbox: %v", err)
		return err
	}

	// Register core functions
	if err := s.registerCoreFunctions(state); err != nil {
		s.debugLog("Failed to register core functions: %v", err)
		return err
	}

	// Register custom modules in the global environment
	s.mu.RLock()
	for name, module := range s.modules {
		s.debugLog("Registering module: %s", name)
		if err := state.PushValue(module); err != nil {
			s.mu.RUnlock()
			s.debugLog("Failed to register module %s: %v", name, err)
			return err
		}
		state.SetGlobal(name)
	}
	s.mu.RUnlock()

	s.debugLog("Sandbox setup complete")
	return nil
}

// registerCoreFunctions registers all built-in functions in the Lua state
func (s *Sandbox) registerCoreFunctions(state *luajit.State) error {
	// Register HTTP functions
	if err := state.RegisterGoFunction("__http_request", httpRequest); err != nil {
		return err
	}

	// Register utility functions
	if err := state.RegisterGoFunction("__generate_token", generateToken); err != nil {
		return err
	}

	// Additional registrations can be added here

	return nil
}

// Execute runs a Lua script in the sandbox with the given context
func (s *Sandbox) Execute(state *luajit.State, bytecode []byte, ctx *Context) (*Response, error) {
	// Create a response object
	response := NewResponse()

	// Load bytecode
	if err := state.LoadBytecode(bytecode, "script"); err != nil {
		ReleaseResponse(response)
		return nil, fmt.Errorf("failed to load script: %w", err)
	}

	// Add session data to context
	contextWithSession := make(map[string]any)
	maps.Copy(contextWithSession, ctx.Values)

	// Pass session data through context
	if ctx.SessionID != "" {
		contextWithSession["session_id"] = ctx.SessionID
		contextWithSession["session_data"] = ctx.SessionData
	}

	// Set up context values for execution
	if err := state.PushTable(contextWithSession); err != nil {
		ReleaseResponse(response)
		return nil, err
	}

	// Get the execution function
	state.GetGlobal("__execute_script")
	if !state.IsFunction(-1) {
		state.Pop(1)
		ReleaseResponse(response)
		return nil, ErrSandboxNotInitialized
	}

	// Push function and bytecode
	state.PushCopy(-2) // Bytecode
	state.PushCopy(-2) // Context
	state.Remove(-4)   // Remove bytecode duplicate
	state.Remove(-3)   // Remove context duplicate

	// Execute with 2 args, 1 result
	if err := state.Call(2, 1); err != nil {
		ReleaseResponse(response)
		return nil, fmt.Errorf("script execution failed: %w", err)
	}

	// Set response body from result
	body, err := state.ToValue(-1)
	if err == nil {
		response.Body = body
	}
	state.Pop(1)

	extractHTTPResponseData(state, response)

	extractSessionData(state, response)

	return response, nil
}

// extractResponseData pulls response info from the Lua state
func extractHTTPResponseData(state *luajit.State, response *Response) {
	state.GetGlobal("__http_responses")
	if !state.IsTable(-1) {
		state.Pop(1)
		return
	}

	state.PushNumber(1)
	state.GetTable(-2)
	if !state.IsTable(-1) {
		state.Pop(2)
		return
	}

	// Extract status
	state.GetField(-1, "status")
	if state.IsNumber(-1) {
		response.Status = int(state.ToNumber(-1))
	}
	state.Pop(1)

	// Extract headers
	state.GetField(-1, "headers")
	if state.IsTable(-1) {
		state.PushNil() // Start iteration
		for state.Next(-2) {
			if state.IsString(-2) && state.IsString(-1) {
				key := state.ToString(-2)
				value := state.ToString(-1)
				response.Headers[key] = value
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Extract cookies
	state.GetField(-1, "cookies")
	if state.IsTable(-1) {
		length := state.GetTableLength(-1)
		for i := 1; i <= length; i++ {
			state.PushNumber(float64(i))
			state.GetTable(-2)

			if state.IsTable(-1) {
				extractCookie(state, response)
			}
			state.Pop(1)
		}
	}
	state.Pop(1)

	// Extract metadata
	state.GetField(-1, "metadata")
	if state.IsTable(-1) {
		table, err := state.ToTable(-1)
		if err == nil {
			for k, v := range table {
				response.Metadata[k] = v
			}
		}
	}
	state.Pop(1)

	// Clean up
	state.Pop(2)
}

// extractCookie pulls cookie data from the current table on the stack
func extractCookie(state *luajit.State, response *Response) {
	cookie := fasthttp.AcquireCookie()

	// Get name (required)
	state.GetField(-1, "name")
	if !state.IsString(-1) {
		state.Pop(1)
		fasthttp.ReleaseCookie(cookie)
		return
	}
	cookie.SetKey(state.ToString(-1))
	state.Pop(1)

	// Get value
	state.GetField(-1, "value")
	if state.IsString(-1) {
		cookie.SetValue(state.ToString(-1))
	}
	state.Pop(1)

	// Get path
	state.GetField(-1, "path")
	if state.IsString(-1) {
		cookie.SetPath(state.ToString(-1))
	} else {
		cookie.SetPath("/") // Default
	}
	state.Pop(1)

	// Get domain
	state.GetField(-1, "domain")
	if state.IsString(-1) {
		cookie.SetDomain(state.ToString(-1))
	}
	state.Pop(1)

	// Get other parameters
	state.GetField(-1, "http_only")
	if state.IsBoolean(-1) {
		cookie.SetHTTPOnly(state.ToBoolean(-1))
	}
	state.Pop(1)

	state.GetField(-1, "secure")
	if state.IsBoolean(-1) {
		cookie.SetSecure(state.ToBoolean(-1))
	}
	state.Pop(1)

	state.GetField(-1, "max_age")
	if state.IsNumber(-1) {
		cookie.SetMaxAge(int(state.ToNumber(-1)))
	}
	state.Pop(1)

	response.Cookies = append(response.Cookies, cookie)
}

// Extract session data if modified
func extractSessionData(state *luajit.State, response *Response) {
	logger.Debug("extractSessionData: Starting extraction")

	// Get HTTP response table
	state.GetGlobal("__http_responses")
	if !state.IsTable(-1) {
		logger.Debug("extractSessionData: __http_responses is not a table")
		state.Pop(1)
		return
	}

	// Get first response
	state.PushNumber(1)
	state.GetTable(-2)
	if !state.IsTable(-1) {
		logger.Debug("extractSessionData: __http_responses[1] is not a table")
		state.Pop(2)
		return
	}

	// Check session_modified flag
	state.GetField(-1, "session_modified")
	if !state.IsBoolean(-1) || !state.ToBoolean(-1) {
		logger.Debug("extractSessionData: session_modified is not true")
		state.Pop(3)
		return
	}
	logger.Debug("extractSessionData: Found session_modified=true")
	state.Pop(1)

	// Get session ID
	state.GetField(-1, "session_id")
	if state.IsString(-1) {
		response.SessionID = state.ToString(-1)
		logger.Debug("extractSessionData: Found session ID: %s", response.SessionID)
	} else {
		logger.Debug("extractSessionData: session_id not found or not a string")
	}
	state.Pop(1)

	// Get session data
	state.GetField(-1, "session_data")
	if state.IsTable(-1) {
		logger.Debug("extractSessionData: Found session_data table")
		sessionData, err := state.ToTable(-1)
		if err == nil {
			logger.Debug("extractSessionData: Converted session data, size=%d", len(sessionData))
			for k, v := range sessionData {
				response.SessionData[k] = v
				logger.Debug("extractSessionData: Added session key=%s, value=%v", k, v)
			}
			response.SessionModified = true
		} else {
			logger.Debug("extractSessionData: Failed to convert session data: %v", err)
		}
	} else {
		logger.Debug("extractSessionData: session_data not found or not a table")
	}
	state.Pop(1)

	// Clean up stack
	state.Pop(2)
	logger.Debug("extractSessionData: Finished extraction, modified=%v", response.SessionModified)
}