package daemon

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/AlexanderGrooff/spage-protobuf/spage/core"
	"github.com/AlexanderGrooff/spage/pkg/common"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// Client represents a gRPC client for communicating with the Spage daemon
type Client struct {
	conn   *grpc.ClientConn
	client core.SpageExecutionClient

	// Configuration
	endpoint string
	playID   string
	timeout  time.Duration

	// Connection management
	mu        sync.RWMutex
	connected bool

	// Progress streaming
	progressStream     core.SpageExecution_StreamTaskProgressClient
	playProgressStream core.SpageExecution_StreamPlayProgressClient
	streamMu           sync.RWMutex

	// Context for graceful shutdown
	ctx    context.Context
	cancel context.CancelFunc
}

// Config holds the client configuration
type Config struct {
	Endpoint string
	PlayID   string
	Timeout  time.Duration
}

// NewClient creates a new gRPC client for daemon communication
func NewClient(cfg *Config) (*Client, error) {
	if cfg.Endpoint == "" {
		return nil, fmt.Errorf("endpoint is required")
	}

	if cfg.Timeout == 0 {
		cfg.Timeout = 3 * time.Second
	}

	ctx, cancel := context.WithCancel(context.Background())

	client := &Client{
		endpoint: cfg.Endpoint,
		playID:   cfg.PlayID,
		timeout:  cfg.Timeout,
		ctx:      ctx,
		cancel:   cancel,
	}

	return client, nil
}

// Connect establishes a connection to the daemon
func (c *Client) Connect() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.connected {
		return nil
	}

	conn, err := grpc.NewClient(
		c.endpoint,
		grpc.WithTransportCredentials(insecure.NewCredentials()),
	)
	if err != nil {
		// Check if it's a connection error
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available at %s: %w", c.endpoint, err)
		}
		return fmt.Errorf("failed to connect to daemon: %w", err)
	}

	c.conn = conn
	c.client = core.NewSpageExecutionClient(conn)
	c.connected = true

	return nil
}

// Disconnect closes the connection to the daemon
func (c *Client) Disconnect() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if !c.connected {
		return nil
	}

	// Close progress streams gracefully
	if err := c.CloseStreamsGracefully(); err != nil {
		common.LogWarn("Failed to close streams gracefully", map[string]interface{}{
			"error": err.Error(),
		})
	}

	// Close connection
	if c.conn != nil {
		if err := c.conn.Close(); err != nil {
			return fmt.Errorf("failed to close connection: %w", err)
		}
		c.conn = nil
	}

	c.connected = false
	return nil
}

// RegisterPlayStart registers the current play with the daemon via streaming
func (c *Client) RegisterPlayStart(playbook, inventory string, variables map[string]string, executor string) error {
	if err := c.ensureConnected(); err != nil {
		// Check if it's a connection error
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available: %w", err)
		}
		return err
	}

	// Ensure play progress stream is established
	if err := c.ensurePlayProgressStream(); err != nil {
		return fmt.Errorf("failed to ensure play progress stream: %w", err)
	}

	update := &core.PlayProgressUpdate{
		PlayId:    c.playID,
		Status:    core.PlayStatus_PLAY_STATUS_PENDING,
		Output:    "Play started",
		Timestamp: timestamppb.Now(),
	}

	c.streamMu.RLock()
	defer c.streamMu.RUnlock()

	if c.playProgressStream != nil {
		// Try to send the update with retry logic
		var err error
		for retries := 0; retries < 3; retries++ {
			err = c.playProgressStream.Send(update)
			if err == nil {
				break
			}

			// If it's a connection error, try to reconnect
			if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
				c.streamMu.RUnlock()
				c.streamMu.Lock()
				c.playProgressStream = nil // Reset stream to force reconnection
				c.streamMu.Unlock()
				c.streamMu.RLock()

				// Try to reestablish connection
				if reconnectErr := c.ensurePlayProgressStream(); reconnectErr != nil {
					return fmt.Errorf("failed to reestablish play progress stream: %w", reconnectErr)
				}
				continue
			}

			// For other errors, don't retry
			break
		}

		if err != nil {
			return fmt.Errorf("failed to send play start update: %w", err)
		}
	}

	return nil
}

func (c *Client) RegisterPlayCompletion() error {
	if err := c.ensureConnected(); err != nil {
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available: %w", err)
		}
		return err
	}

	// Ensure play progress stream is established
	if err := c.ensurePlayProgressStream(); err != nil {
		return fmt.Errorf("failed to ensure play progress stream: %w", err)
	}

	update := &core.PlayProgressUpdate{
		PlayId:    c.playID,
		Status:    core.PlayStatus_PLAY_STATUS_COMPLETED,
		Output:    "Play completed",
		Timestamp: timestamppb.Now(),
	}

	common.LogDebug("Registering play completion", map[string]interface{}{
		"play_id": c.playID,
	})

	c.streamMu.RLock()
	defer c.streamMu.RUnlock()

	if c.playProgressStream != nil {
		// Try to send the update with retry logic
		var err error
		for retries := 0; retries < 3; retries++ {
			err = c.playProgressStream.Send(update)
			if err == nil {
				break
			}

			// If it's a connection error, try to reconnect
			if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
				c.streamMu.RUnlock()
				c.streamMu.Lock()
				c.playProgressStream = nil // Reset stream to force reconnection
				c.streamMu.Unlock()
				c.streamMu.RLock()

				// Try to reestablish connection
				if reconnectErr := c.ensurePlayProgressStream(); reconnectErr != nil {
					return fmt.Errorf("failed to reestablish play progress stream: %w", reconnectErr)
				}
				continue
			}

			// For other errors, don't retry
			break
		}

		if err != nil {
			return fmt.Errorf("failed to send play completion update: %w", err)
		}
	}

	return nil
}

// RegisterPlayError registers a play error with the daemon via streaming
func (c *Client) RegisterPlayError(err error) error {
	if err := c.ensureConnected(); err != nil {
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available: %w", err)
		}
		return err
	}

	// Ensure play progress stream is established
	if err := c.ensurePlayProgressStream(); err != nil {
		return fmt.Errorf("failed to ensure play progress stream: %w", err)
	}

	update := &core.PlayProgressUpdate{
		PlayId:    c.playID,
		Status:    core.PlayStatus_PLAY_STATUS_FAILED,
		Error:     err.Error(),
		Output:    "Play failed",
		Timestamp: timestamppb.Now(),
	}

	common.LogDebug("Registering play error", map[string]interface{}{
		"play_id": c.playID,
	})

	c.streamMu.RLock()
	defer c.streamMu.RUnlock()

	if c.playProgressStream != nil {
		// Try to send the update with retry logic
		var err error
		for retries := 0; retries < 3; retries++ {
			err = c.playProgressStream.Send(update)
			if err == nil {
				break
			}

			// If it's a connection error, try to reconnect
			if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
				c.streamMu.RUnlock()
				c.streamMu.Lock()
				c.playProgressStream = nil // Reset stream to force reconnection
				c.streamMu.Unlock()
				c.streamMu.RLock()

				// Try to reestablish connection
				if reconnectErr := c.ensurePlayProgressStream(); reconnectErr != nil {
					return fmt.Errorf("failed to reestablish play progress stream: %w", reconnectErr)
				}
				continue
			}

			// For other errors, don't retry
			break
		}

		if err != nil {
			return fmt.Errorf("failed to send play error update: %w", err)
		}
	}

	return nil
}

// UpdateTaskResult sends a task result update to the daemon
func (c *Client) UpdateTaskResult(taskResult *core.TaskResult) error {
	if c == nil {
		return fmt.Errorf("cannot update task result: client is nil")
	}
	common.LogInfo("Updating task result", map[string]interface{}{
		"task_id": taskResult.TaskId,
		"status":  taskResult.Status,
	})

	// Check if we're connected first
	c.mu.RLock()
	connected := c.connected
	c.mu.RUnlock()

	if !connected {
		if err := c.Connect(); err != nil {
			return fmt.Errorf("failed to connect to daemon: %w", err)
		}
	}

	// Ensure progress stream is established
	if err := c.ensureProgressStream(); err != nil {
		return fmt.Errorf("failed to ensure progress stream: %w", err)
	}

	update := &core.TaskProgressUpdate{
		PlayId:    c.playID,
		TaskId:    int32(taskResult.TaskId),
		Result:    taskResult,
		Timestamp: timestamppb.Now(),
	}

	c.streamMu.RLock()
	defer c.streamMu.RUnlock()

	if c.progressStream != nil {
		// Try to send the update with retry logic
		var err error
		for retries := 0; retries < 3; retries++ {
			err = c.progressStream.Send(update)
			if err == nil {
				break
			}

			// If it's a connection error, try to reconnect
			if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
				c.streamMu.RUnlock()
				c.streamMu.Lock()
				c.progressStream = nil // Reset stream to force reconnection
				c.streamMu.Unlock()
				c.streamMu.RLock()

				// Try to reestablish connection
				if reconnectErr := c.ensureProgressStream(); reconnectErr != nil {
					return fmt.Errorf("failed to reestablish progress stream: %w", reconnectErr)
				}
				continue
			}

			// For other errors, don't retry
			break
		}

		if err != nil {
			return fmt.Errorf("failed to send task result update: %w", err)
		}
	}

	return nil
}

// ensureConnected ensures the client is connected to the daemon
func (c *Client) ensureConnected() error {
	if c == nil {
		return fmt.Errorf("daemon client is nil")
	}

	c.mu.RLock()
	connected := c.connected
	c.mu.RUnlock()

	if !connected {
		return c.Connect()
	}

	return nil
}

// ensureProgressStream ensures the progress stream is established
func (c *Client) ensureProgressStream() error {
	if c == nil {
		return fmt.Errorf("daemon client is nil")
	}

	if err := c.ensureConnected(); err != nil {
		return err
	}

	c.streamMu.Lock()
	defer c.streamMu.Unlock()

	if c.progressStream != nil {
		return nil
	}

	// Create progress stream
	stream, err := c.client.StreamTaskProgress(c.ctx)
	if err != nil {
		// Check if it's a connection error
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available: %w", err)
		}
		return fmt.Errorf("failed to create progress stream: %w", err)
	}

	c.progressStream = stream

	// Start receiving updates in background
	go c.receiveProgressUpdates()

	return nil
}

// ensurePlayProgressStream ensures the play progress stream is established
func (c *Client) ensurePlayProgressStream() error {
	if c == nil {
		return fmt.Errorf("daemon client is nil")
	}

	if err := c.ensureConnected(); err != nil {
		return err
	}

	c.streamMu.Lock()
	defer c.streamMu.Unlock()

	if c.playProgressStream != nil {
		return nil
	}

	// Create play progress stream
	stream, err := c.client.StreamPlayProgress(c.ctx)
	if err != nil {
		// Check if it's a connection error
		if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
			return fmt.Errorf("daemon not available: %w", err)
		}
		return fmt.Errorf("failed to create play progress stream: %w", err)
	}

	c.playProgressStream = stream

	return nil
}

// receiveProgressUpdates receives progress updates from the daemon
func (c *Client) receiveProgressUpdates() {
	for {
		// Check if stream is still valid
		c.streamMu.RLock()
		stream := c.progressStream
		c.streamMu.RUnlock()

		if stream == nil {
			// Stream was closed or reset, exit gracefully
			return
		}

		update, err := stream.Recv()
		if err != nil {
			// Check if it's a graceful close
			if status.Code(err) == codes.Canceled {
				return
			}

			// For connection errors, exit gracefully without logging
			if status.Code(err) == codes.Unavailable || status.Code(err) == codes.DeadlineExceeded {
				return
			}

			// For other errors, log and exit
			fmt.Printf("Error receiving progress update: %v\n", err)
			return
		}

		// Handle incoming progress updates (if needed)
		if update.Result != nil {
			fmt.Printf("Received task result: %d - %s\n", update.TaskId, update.Result.Status)
		} else {
			fmt.Printf("Received progress update: %d\n", update.TaskId)
		}
	}
}

// Close closes the client and cleans up resources
func (c *Client) Close() error {
	c.cancel()
	return c.Disconnect()
}

// CloseStreamsGracefully closes the streams gracefully, waiting for any pending data
func (c *Client) CloseStreamsGracefully() error {
	c.streamMu.Lock()
	defer c.streamMu.Unlock()

	// Close progress stream gracefully
	if c.progressStream != nil {
		// CloseSend() closes the sending side of the stream
		if err := c.progressStream.CloseSend(); err != nil {
			common.LogWarn("Failed to close progress stream send", map[string]interface{}{
				"error": err.Error(),
			})
		}

		c.progressStream = nil
	}

	// Close play progress stream gracefully
	if c.playProgressStream != nil {
		if err := c.playProgressStream.CloseSend(); err != nil {
			common.LogWarn("Failed to close play progress stream send", map[string]interface{}{
				"error": err.Error(),
			})
		}

		c.playProgressStream = nil
	}

	return nil
}

// WaitForStreamsToFinish waits for streams to finish processing any pending data
func (c *Client) WaitForStreamsToFinish(timeout time.Duration) error {
	deadline := time.Now().Add(timeout)

	for time.Now().Before(deadline) {
		c.streamMu.RLock()
		hasProgressStream := c.progressStream != nil
		hasPlayProgressStream := c.playProgressStream != nil
		c.streamMu.RUnlock()

		// If no streams are active, we're done
		if !hasProgressStream && !hasPlayProgressStream {
			return nil
		}

		// Wait a bit before checking again
		time.Sleep(50 * time.Millisecond)
	}

	return fmt.Errorf("timeout waiting for streams to finish after %v", timeout)
}

// IsConnected returns whether the client is connected to the daemon
func (c *Client) IsConnected() bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.connected
}
