package handlers

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"
	"time"

	"github.com/go-redis/redis/v8"
	"github.com/gorilla/websocket"
	redis_cli "github.com/suno-ai/glockenspiel/studio_api_flute/internal/redis"
	"github.com/suno-ai/glockenspiel/studio_api_flute/internal/studio_api"
)

type WebSocketRequest struct {
	Type   string                 `json:"type"`
	Method string                 `json:"method"`
	Path   string                 `json:"path"`
	Params map[string]interface{} `json:"params"`
}

type GenerateV2Response struct {
	ID                *string   `json:"id,omitempty"`
	Clips             *[]Clip   `json:"clips,omitempty"`
	Metadata          *Metadata `json:"metadata,omitempty"`
	MajorModelVersion *string   `json:"major_model_version,omitempty"`
	Status            *string   `json:"status,omitempty"`
	CreatedAt         *string   `json:"created_at,omitempty"`
	BatchSize         *int32    `json:"batch_size,omitempty"`
	Message           *string   `json:"message,omitempty"`
}

type Clip struct {
	ID                *string   `json:"id,omitempty"`
	VideoURL          *string   `json:"video_url,omitempty"`
	AudioURL          *string   `json:"audio_url,omitempty"`
	ImageURL          *string   `json:"image_url,omitempty"`
	ImageLargeURL     *string   `json:"image_large_url,omitempty"`
	IsVideoPending    *bool     `json:"is_video_pending,omitempty"`
	MajorModelVersion *string   `json:"major_model_version,omitempty"`
	ModelName         *string   `json:"model_name,omitempty"`
	Metadata          *Metadata `json:"metadata,omitempty"`
	IsLiked           *bool     `json:"is_liked,omitempty"`
	UserID            *string   `json:"user_id,omitempty"`
	DisplayName       *string   `json:"display_name,omitempty"`
	Handle            *string   `json:"handle,omitempty"`
	IsHandleUpdated   *bool     `json:"is_handle_updated,omitempty"`
	AvatarImageURL    *string   `json:"avatar_image_url,omitempty"`
	IsTrashed         *bool     `json:"is_trashed,omitempty"`
	Reaction          *Reaction `json:"reaction,omitempty"`
	CreatedAt         *string   `json:"created_at,omitempty"`
	Status            *string   `json:"status,omitempty"`
	Title             *string   `json:"title,omitempty"`
	PlayCount         *int32    `json:"play_count,omitempty"`
	UpvoteCount       *int32    `json:"upvote_count,omitempty"`
	IsPublic          *bool     `json:"is_public,omitempty"`
}

type HistoryItem struct {
	ID         *string  `json:"id,omitempty"`
	ContinueAt *float64 `json:"continue_at,omitempty"`
	Type       *string  `json:"type,omitempty"`
	Source     *string  `json:"source,omitempty"`
	Infill     *bool    `json:"infill,omitempty"`
}

type Metadata struct {
	Tags                      *string           `json:"tags,omitempty"`
	NegativeTags              *string           `json:"negative_tags,omitempty"`
	Prompt                    *string           `json:"prompt,omitempty"`
	GptDescriptionPrompt      *string           `json:"gpt_description_prompt,omitempty"`
	History                   *[]HistoryItem    `json:"history,omitempty"`
	ConcatHistory             *[]HistoryItem    `json:"concat_history,omitempty"`
	StemFromID                *string           `json:"stem_from_id,omitempty"`
	Type                      *string           `json:"type,omitempty"`
	Duration                  *float64          `json:"duration,omitempty"`
	RefundCredits             *bool             `json:"refund_credits,omitempty"`
	Stream                    *bool             `json:"stream,omitempty"`
	Infill                    *bool             `json:"infill,omitempty"`
	HasVocal                  *bool             `json:"has_vocal,omitempty"`
	IsAudioUploadTosAccepted  *bool             `json:"is_audio_upload_tos_accepted,omitempty"`
	ErrorType                 *string           `json:"error_type,omitempty"`
	ErrorMessage              *string           `json:"error_message,omitempty"`
	Configurations            *map[string]int32 `json:"configurations,omitempty"`
	ArtistClipID              *string           `json:"artist_clip_id,omitempty"`
	CoverClipID               *string           `json:"cover_clip_id,omitempty"`
	PersonaID                 *string           `json:"persona_id,omitempty"`
	VideoToSongVideoUploadID  *string           `json:"video_to_song_video_upload_id,omitempty"`
	VideoToSongVideoUploadURL *string           `json:"video_to_song_video_upload_url,omitempty"`
	VideoToSongVideoOutputURL *string           `json:"video_to_song_video_output_url,omitempty"`
	ImageToSongImageIDs       *[]string         `json:"image_to_song_image_ids,omitempty"`
	ImageToSongImageURLs      *[]string         `json:"image_to_song_image_urls,omitempty"`
	IsImageToSong             *bool             `json:"is_image_to_song,omitempty"`
	IsSunoShort               *bool             `json:"is_suno_short,omitempty"`
	Task                      *string           `json:"task,omitempty"`
	EditSessionID             *string           `json:"edit_session_id,omitempty"`
}

type Reaction struct {
	PlayCount      *int32  `json:"play_count,omitempty"`
	SkipCount      *int32  `json:"skip_count,omitempty"`
	Flagged        *bool   `json:"flagged,omitempty"`
	FlaggedReason  *string `json:"flagged_reason,omitempty"`
	FeedbackReason *string `json:"feedback_reason,omitempty"`
	ReactionType   *string `json:"reaction_type,omitempty"`
	Clip           *string `json:"clip,omitempty"`
	UpdatedAt      *string `json:"updated_at,omitempty"`
}

type TextPayload struct {
	RequestID string `json:"request_id"`
	ID        string `json:"id"`
	Title     string `json:"title"`
	Type      string `json:"type"`
	Text      string `json:"text"`
	Tags      string `json:"tags"`
	Lang      string `json:"lang"`
}

type ExperimentPayload struct {
	RequestID         string   `json:"request_id"`
	ID                string   `json:"id"`
	Type              string   `json:"type"`
	Experiment        []string `json:"experiment"`
	ExperimentVersion string   `json:"experiment_version"`
	ModelName         string   `json:"model_name"`
}

type StreamingPayload struct {
	RequestID string `json:"request_id"`
	ID        string `json:"id"`
	Type      string `json:"type"`
}

type ImagePayload struct {
	RequestID string `json:"request_id"`
	ID        string `json:"id"`
	Type      string `json:"type"`
	ImageID   string `json:"image_id"`
}

type AudioPayload struct {
	RequestID   string    `json:"request_id"`
	ID          string    `json:"id"`
	Model       string    `json:"model"`
	NAudios     int       `json:"n_audios"`
	OK          bool      `json:"ok"`
	GenDuration float64   `json:"gen_duration"`
	IDs         []string  `json:"ids"`
	Durations   []float64 `json:"durations"`
}

type ErrorPayload struct {
	RequestID    string `json:"request_id"`
	ID           string `json:"id"`
	Type         string `json:"type"`
	ErrorType    string `json:"error_type"`
	ErrorMessage string `json:"error_message"`
}

type PayloadResult struct {
	Payload interface{}
	Type    string
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  8192,
	WriteBufferSize: 8192,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type WebSocketPayload struct {
	Type string   `json:"type"`
	Data []string `json:"data"`
}

// ExtractAuthTokenFromHeader extracts the authorization token from the request headers
func ExtractAuthTokenFromHeader(r *http.Request) (string, error) {
	authHeader := r.Header.Get("Authorization")
	authHeader = strings.Replace(authHeader, "Bearer ", "", 1)
	if authHeader == "" {
		return "", fmt.Errorf("missing authorization header")
	}

	return authHeader, nil
}

// ExtractAuthTokenFromQuery extracts the authorization token from the query parameters
func ExtractAuthTokenFromQuery(r *http.Request) (string, error) {
	queryParams := r.URL.Query()
	authToken := queryParams.Get("token")
	fmt.Println("authToken:", authToken)
	if authToken == "" {
		return "", fmt.Errorf("missing token in query parameters")
	}

	return authToken, nil
}

// var songRegexp = regexp.MustCompile(`^song:(next|last):([0-9a-fA-F-]{36})$`)
var songRegexp = regexp.MustCompile(`^song:(last):([0-9a-fA-F-]{36})$`)

func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Failed to upgrade to WebSocket: %v", err)
		return
	}
	defer ws.Close()

	authToken, err := ExtractAuthTokenFromHeader(r)
	fmt.Println("authToken:", authToken)
	if err != nil {
		authToken, err = ExtractAuthTokenFromQuery(r)
		if err != nil {
			log.Printf("Error extracting auth token: %v", err)
			return
		}
	}

	messageChan := make(chan []byte)
	errorChan := make(chan error)

	// Channel to add new response channels
	newResponseChan := make(chan (<-chan *redis.Message))

	// Goroutine to receive messages
	go func() {
		for {
			_, msg, err := ws.ReadMessage()
			if err != nil {
				errorChan <- err
				return
			}
			messageChan <- msg
		}
	}()

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

	// Goroutine to process and send messages
	go func() {
		for {
			select {
			case msg := <-messageChan:
				var request WebSocketRequest
				if err := json.Unmarshal(msg, &request); err != nil {
					log.Printf("Error unmarshalling request: %v", err)
					continue
				}

				if request.Type == "studio" {
					responseMessage, err := studio_api.SendStudioApiRequest(authToken, request.Method, request.Path, request.Params)
					if err != nil {
						log.Printf("Error sending request to api server: %v", err)
						continue
					}

					// To be removed after adding request_id to studio_api
					var resp *GenerateV2Response
					if err := json.Unmarshal([]byte(responseMessage), &resp); err != nil {
						log.Printf("Error unmarshalling response: %v", err)
						return
					}

					for _, clip := range *resp.Clips {
						ch := redis_cli.GetRedisClient().SubscribeToChannel(*clip.ID)
						newResponseChan <- ch
					}

					// Subscribe to the Redis channel for the response
					responseChan := redis_cli.GetRedisClient().SubscribeToChannel(*resp.ID)
					newResponseChan <- responseChan

					if err := ws.WriteMessage(websocket.TextMessage, []byte(responseMessage)); err != nil {
						log.Printf("Error writing message: %v", err)
						return
					}
				} else if request.Type == "stream" {
					// StreamSet Name: streams_set
					// SongList in a stream: {stream_name}_song_list
					// Stream_audio_chunk: {stream_name}_audio_chunk
					// song_name: https://cdn1.suno.ai/dc55f8d0-2f67-4f45-8b5d-2e6677c0da1d.mp3
					// Create a context with cancel function

					if request.Path != "suno_stream" {
						log.Printf("Only suno_stream is supported for now, %v is not supported", request.Path)
						continue
					}

					if request.Method == "subscribe" {
						// Subscribe to the suno_stream:events channel
						pubsub := redis_cli.GetRedisClient().Subscribe("suno_stream:events")
						defer pubsub.Close()

						// Fetch the song list
						songList, err := redis_cli.GetRedisClient().FetchRedisList("suno_stream:song_list")
						if err != nil {
							log.Printf("Error fetching song list: %v", err)
							continue
						}

						// Wrap the song list in a JSON payload with an additional field "type": "song_list"
						songListPayload := WebSocketPayload{
							Type: "song_list",
							Data: songList,
						}

						// Send the song list to the WebSocket client
						songListMessage, err := json.Marshal(songListPayload)
						if err != nil {
							log.Printf("Error marshalling song list: %v", err)
							continue
						}
						if err := ws.WriteMessage(websocket.TextMessage, songListMessage); err != nil {
							log.Printf("Error writing song list message: %v", err)
							return
						}

						// Send the comments to the WebSocket client
						comments, err := redis_cli.GetRedisClient().FetchRedisList("suno_stream:comments")
						if err != nil {
							log.Printf("Error fetching comments: %v", err)
							continue
						}

						commentsPayload := WebSocketPayload{
							Type: "comments",
							Data: comments,
						}
						commentsMessage, err := json.Marshal(commentsPayload)
						if err != nil {
							log.Printf("Error marshalling comments: %v", err)
							continue
						}
						if err := ws.WriteMessage(websocket.TextMessage, commentsMessage); err != nil {
							log.Printf("Error writing comments message: %v", err)
							return
						}

						// Listen for messages from the subscribed channel
						go func() {
							for {
								msg, err := pubsub.ReceiveMessage(ctx)
								if err != nil {
									log.Printf("Error receiving message: %v", err)
									return
								}

								eventPayload := WebSocketPayload{
									Type: "event",
									Data: []string{msg.Payload},
								}
								eventMessage, err := json.Marshal(eventPayload)
								if err != nil {
									log.Printf("Error unmarshalling event payload: %v", err)
									continue
								}

								if err := ws.WriteMessage(websocket.TextMessage, eventMessage); err != nil {
									log.Printf("Error writing message: %v", err)
									return
								}
							}
						}()
					} else if request.Method == "comment" {
						text, ok := request.Params["text"]
						if !ok {
							log.Printf("Missing text in params")
							continue
						}
						textStr := text.(string)

						// Check if textStr is in the format "song:next:e1afffe6-a566-4204-bc7f-1ddf273937e5" or "song:last:e1afffe6-a566-4204-bc7f-1ddf273937e5"
						matches := songRegexp.FindStringSubmatch(textStr)
						fmt.Println("matches:", matches)
						if matches != nil {
							action := matches[1] // "next" or "last"
							uuid := matches[2]   // UUID part
							log.Printf("Text is in the song addition format: %s, action: %s, uuid: %s", textStr, action, uuid)
							redis_cli.GetRedisClient().AddToRedisList("suno_stream:song_list", "https://cdn1.suno.ai/"+uuid+".mp3")

							// Fetch the song list
							songList, err := redis_cli.GetRedisClient().FetchRedisList("suno_stream:song_list")
							if err != nil {
								log.Printf("Error fetching song list: %v", err)
								continue
							}

							// Wrap the song list in a JSON payload with an additional field "type": "song_list"
							songListPayload := WebSocketPayload{
								Type: "song_list",
								Data: songList,
							}

							// Send the song list to the WebSocket client
							songListMessage, err := json.Marshal(songListPayload)
							if err != nil {
								log.Printf("Error marshalling song list: %v", err)
								continue
							}
							if err := ws.WriteMessage(websocket.TextMessage, songListMessage); err != nil {
								log.Printf("Error writing song list message: %v", err)
								return
							}
						} else {
							if err := redis_cli.GetRedisClient().PublishMessage("suno_stream:events", "comment:"+textStr); err != nil {
								log.Printf("Error publishing message: %v", err)
								continue
							}
							redis_cli.GetRedisClient().AddToRedisList("suno_stream:comments", textStr)
							log.Printf("Text is not in the correct format: %s", textStr)
						}
					} else {
						log.Printf("Invalid method: %v", request.Method)
						continue
					}
				} else {
					log.Printf("Invalid type: %v", request.Type)
					continue
				}

			case err := <-errorChan:
				log.Printf("Error: %v", err)
				return
			}
		}
	}()

	// Merge all response channels into mergedChannel dynamically
	mergedChannel := redis_cli.MergeChannelsDynamically(newResponseChan)

	go func() {
		for {
			select {
			case msg := <-mergedChannel:
				if msg == nil {
					log.Printf("Received nil message")
					continue
				}
				fmt.Println("Received message from channel:", msg.Channel, "Message:", msg.Payload)
				if err := ws.WriteMessage(websocket.TextMessage, []byte(msg.Payload)); err != nil {
					log.Printf("Error writing message: %v", err)
					return
				}
			default:
				time.Sleep(20 * time.Millisecond)
			}
		}
	}()

	// Block until an error occurs
	err = <-errorChan
	log.Printf("WebSocket error: %v", err)
}