package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"
)

// Tetrimino types
const (
	EMPTY = iota
	I_PIECE
	O_PIECE
	T_PIECE
	S_PIECE
	Z_PIECE
	J_PIECE
	L_PIECE
)

// Tetrimino shapes (4 rotations each)
var tetriminoShapes = map[int][4][4][4]int{
	I_PIECE: {
		{{0, 0, 0, 0}, {1, 1, 1, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 0, 1, 0}, {0, 0, 1, 0}, {0, 0, 1, 0}, {0, 0, 1, 0}},
		{{0, 0, 0, 0}, {0, 0, 0, 0}, {1, 1, 1, 1}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {0, 1, 0, 0}, {0, 1, 0, 0}, {0, 1, 0, 0}},
	},
	O_PIECE: {
		{{0, 1, 1, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 1, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 1, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 1, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
	},
	T_PIECE: {
		{{0, 1, 0, 0}, {1, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {0, 1, 1, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
		{{0, 0, 0, 0}, {1, 1, 1, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {1, 1, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
	},
	S_PIECE: {
		{{0, 1, 1, 0}, {1, 1, 0, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {0, 1, 1, 0}, {0, 0, 1, 0}, {0, 0, 0, 0}},
		{{0, 0, 0, 0}, {0, 1, 1, 0}, {1, 1, 0, 0}, {0, 0, 0, 0}},
		{{1, 0, 0, 0}, {1, 1, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
	},
	Z_PIECE: {
		{{1, 1, 0, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 0, 1, 0}, {0, 1, 1, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
		{{0, 0, 0, 0}, {1, 1, 0, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {1, 1, 0, 0}, {1, 0, 0, 0}, {0, 0, 0, 0}},
	},
	J_PIECE: {
		{{1, 0, 0, 0}, {1, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 1, 0}, {0, 1, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
		{{0, 0, 0, 0}, {1, 1, 1, 0}, {0, 0, 1, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {0, 1, 0, 0}, {1, 1, 0, 0}, {0, 0, 0, 0}},
	},
	L_PIECE: {
		{{0, 0, 1, 0}, {1, 1, 1, 0}, {0, 0, 0, 0}, {0, 0, 0, 0}},
		{{0, 1, 0, 0}, {0, 1, 0, 0}, {0, 1, 1, 0}, {0, 0, 0, 0}},
		{{0, 0, 0, 0}, {1, 1, 1, 0}, {1, 0, 0, 0}, {0, 0, 0, 0}},
		{{1, 1, 0, 0}, {0, 1, 0, 0}, {0, 1, 0, 0}, {0, 0, 0, 0}},
	},
}

type Tetrimino struct {
	Type     int `json:"type"`
	X        int `json:"x"`
	Y        int `json:"y"`
	Rotation int `json:"rotation"`
}

type GameState struct {
	Board         [40][10]int `json:"board"`
	CurrentPiece  *Tetrimino  `json:"current_piece"`
	NextPieces    []int       `json:"next_pieces"`
	HoldPiece     *int        `json:"hold_piece"`
	Score         int         `json:"score"`
	Level         int         `json:"level"`
	LinesCleared  int         `json:"lines_cleared"`
	GameRunning   bool        `json:"game_running"`
	GameOver      bool        `json:"game_over"`
	CanHold       bool        `json:"can_hold"`
	LinesForLevel int         `json:"lines_for_level"`
}

type Game struct {
	state    *GameState
	bag      []int
	bagIndex int
	clients  map[chan string]bool
}

func NewGame() *Game {
	g := &Game{
		state: &GameState{
			NextPieces:    make([]int, 6),
			Level:         1,
			GameRunning:   false,
			GameOver:      false,
			CanHold:       true,
			LinesForLevel: 5,
		},
		clients: make(map[chan string]bool),
	}
	g.fillBag()
	g.fillNextPieces()
	return g
}

func (g *Game) fillBag() {
	g.bag = []int{I_PIECE, O_PIECE, T_PIECE, S_PIECE, Z_PIECE, J_PIECE, L_PIECE}
	// Simple shuffle using current time
	for i := len(g.bag) - 1; i > 0; i-- {
		j := int(time.Now().UnixNano()) % (i + 1)
		g.bag[i], g.bag[j] = g.bag[j], g.bag[i]
	}
	g.bagIndex = 0
}

func (g *Game) getNextPiece() int {
	if g.bagIndex >= len(g.bag) {
		g.fillBag()
	}
	piece := g.bag[g.bagIndex]
	g.bagIndex++
	return piece
}

func (g *Game) fillNextPieces() {
	for i := 0; i < 6; i++ {
		g.state.NextPieces[i] = g.getNextPiece()
	}
}

func (g *Game) spawnPiece() {
	pieceType := g.state.NextPieces[0]
	// Shift next pieces
	copy(g.state.NextPieces[:5], g.state.NextPieces[1:])
	g.state.NextPieces[5] = g.getNextPiece()

	// Spawn position based on guidelines
	x := 3 // left-middle for most pieces
	if pieceType == I_PIECE || pieceType == O_PIECE {
		x = 4 // middle for I and O
	}

	g.state.CurrentPiece = &Tetrimino{
		Type:     pieceType,
		X:        x,
		Y:        21,
		Rotation: 0,
	}

	// Drop one space if possible
	if g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y-1, g.state.CurrentPiece.Rotation) {
		g.state.CurrentPiece.Y--
	}

	// Check for game over
	if !g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y, g.state.CurrentPiece.Rotation) {
		g.state.GameOver = true
		g.state.GameRunning = false
	}

	g.state.CanHold = true
}

func (g *Game) canMove(x, y, rotation int) bool {
	shape := tetriminoShapes[g.state.CurrentPiece.Type][rotation]
	for py := 0; py < 4; py++ {
		for px := 0; px < 4; px++ {
			if shape[py][px] == 0 {
				continue
			}
			boardX := x + px
			boardY := y + py
			if boardX < 0 || boardX >= 10 || boardY < 0 {
				return false
			}
			if boardY < 40 && g.state.Board[boardY][boardX] != EMPTY {
				return false
			}
		}
	}
	return true
}

func (g *Game) lockPiece() {
	shape := tetriminoShapes[g.state.CurrentPiece.Type][g.state.CurrentPiece.Rotation]
	for py := 0; py < 4; py++ {
		for px := 0; px < 4; px++ {
			if shape[py][px] == 0 {
				continue
			}
			boardX := g.state.CurrentPiece.X + px
			boardY := g.state.CurrentPiece.Y + py
			if boardY < 40 {
				g.state.Board[boardY][boardX] = g.state.CurrentPiece.Type
			}
		}
	}
	g.clearLines()
	g.spawnPiece()
}

func (g *Game) clearLines() {
	linesCleared := 0
	for y := 0; y < 40; y++ {
		full := true
		for x := 0; x < 10; x++ {
			if g.state.Board[y][x] == EMPTY {
				full = false
				break
			}
		}
		if full {
			// Remove line
			for moveY := y; moveY < 39; moveY++ {
				for x := 0; x < 10; x++ {
					g.state.Board[moveY][x] = g.state.Board[moveY+1][x]
				}
			}
			// Clear top line
			for x := 0; x < 10; x++ {
				g.state.Board[39][x] = EMPTY
			}
			linesCleared++
			y-- // Check same line again
		}
	}

	if linesCleared > 0 {
		g.state.LinesCleared += linesCleared
		g.updateScore(linesCleared)
		g.updateLevel()
	}
}

func (g *Game) updateScore(lines int) {
	baseScore := map[int]int{1: 100, 2: 300, 3: 500, 4: 800}
	if score, ok := baseScore[lines]; ok {
		g.state.Score += score * g.state.Level
	}
}

func (g *Game) updateLevel() {
	requiredLines := g.state.Level * g.state.LinesForLevel
	if g.state.LinesCleared >= requiredLines {
		g.state.Level++
	}
}

func (g *Game) move(direction string) {
	if !g.state.GameRunning || g.state.CurrentPiece == nil {
		return
	}

	newX := g.state.CurrentPiece.X
	newY := g.state.CurrentPiece.Y

	switch direction {
	case "left":
		newX--
	case "right":
		newX++
	case "down":
		newY--
	}

	if g.canMove(newX, newY, g.state.CurrentPiece.Rotation) {
		g.state.CurrentPiece.X = newX
		g.state.CurrentPiece.Y = newY
	}
}

func (g *Game) rotate(direction string) {
	if !g.state.GameRunning || g.state.CurrentPiece == nil {
		return
	}

	newRotation := g.state.CurrentPiece.Rotation
	if direction == "cw" {
		newRotation = (newRotation + 1) % 4
	} else {
		newRotation = (newRotation + 3) % 4
	}

	if g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y, newRotation) {
		g.state.CurrentPiece.Rotation = newRotation
	}
}

func (g *Game) hardDrop() {
	if !g.state.GameRunning || g.state.CurrentPiece == nil {
		return
	}

	for g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y-1, g.state.CurrentPiece.Rotation) {
		g.state.CurrentPiece.Y--
	}
	g.lockPiece()
}

func (g *Game) hold() {
	if !g.state.GameRunning || g.state.CurrentPiece == nil || !g.state.CanHold {
		return
	}

	if g.state.HoldPiece == nil {
		g.state.HoldPiece = &g.state.CurrentPiece.Type
		g.spawnPiece()
	} else {
		temp := *g.state.HoldPiece
		g.state.HoldPiece = &g.state.CurrentPiece.Type
		
		x := 3
		if temp == I_PIECE || temp == O_PIECE {
			x = 4
		}
		
		g.state.CurrentPiece = &Tetrimino{
			Type:     temp,
			X:        x,
			Y:        21,
			Rotation: 0,
		}
		
		if g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y-1, g.state.CurrentPiece.Rotation) {
			g.state.CurrentPiece.Y--
		}
	}
	g.state.CanHold = false
}

func (g *Game) startGame() {
	g.state = &GameState{
		NextPieces:    make([]int, 6),
		Level:         1,
		GameRunning:   true,
		GameOver:      false,
		CanHold:       true,
		LinesForLevel: 5,
	}
	g.fillBag()
	g.fillNextPieces()
	g.spawnPiece()
}

func (g *Game) tick() {
	if !g.state.GameRunning || g.state.CurrentPiece == nil {
		return
	}

	if g.canMove(g.state.CurrentPiece.X, g.state.CurrentPiece.Y-1, g.state.CurrentPiece.Rotation) {
		g.state.CurrentPiece.Y--
	} else {
		g.lockPiece()
	}
}

func (g *Game) broadcast() {
	data, err := json.Marshal(g.state)
	if err != nil {
		log.Printf("Error marshaling game state: %v", err)
		return
	}
	message := fmt.Sprintf("data: %s\n\n", string(data))
	
	// Non-blocking broadcast to all clients
	for client := range g.clients {
		select {
		case client <- message:
			// Message sent successfully
		default:
			// Channel full or client blocked, skip this client
			log.Printf("Dropping message for slow client")
		}
	}
}

var game = NewGame()

func main() {
	// Serve static files (index.html)
	http.Handle("/", http.FileServer(http.Dir("./")))
	
	// API endpoints
	http.HandleFunc("/events", handleSSE)
	http.HandleFunc("/action", handleAction)

	// Game loop
	go func() {
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		
		for {
			select {
			case <-ticker.C:
				if game.state.GameRunning {
					game.tick()
					game.broadcast()
				}
			}
		}
	}()

	fmt.Println("Server starting on :8080")
	fmt.Println("Open http://localhost:8080 in your browser")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

func handleSSE(w http.ResponseWriter, r *http.Request) {
	// Critical SSE headers for nginx compatibility
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no") // Disable nginx buffering
	w.Header().Set("Access-Control-Allow-Origin", "*")

	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "Streaming unsupported", http.StatusInternalServerError)
		return
	}

	client := make(chan string, 10) // Buffered channel
	game.clients[client] = true

	// Send initial state
	data, _ := json.Marshal(game.state)
	fmt.Fprintf(w, "data: %s\n\n", string(data))
	flusher.Flush()

	// Heartbeat to keep connection alive
	ticker := time.NewTicker(15 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case message := <-client:
			fmt.Fprint(w, message)
			flusher.Flush()
		case <-ticker.C:
			// Send comment as heartbeat
			fmt.Fprintf(w, ": heartbeat\n\n")
			flusher.Flush()
		case <-r.Context().Done():
			delete(game.clients, client)
			close(client)
			return
		}
	}
}

func handleAction(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods", "POST")
	w.Header().Set("Access-Control-Allow-Headers", "Content-Type")

	if r.Method == "OPTIONS" {
		return
	}

	action := r.URL.Query().Get("action")
	direction := r.URL.Query().Get("direction")

	switch action {
	case "move":
		game.move(direction)
	case "rotate":
		game.rotate(direction)
	case "hard_drop":
		game.hardDrop()
	case "hold":
		game.hold()
	case "start":
		game.startGame()
	}

	game.broadcast()
	w.WriteHeader(http.StatusOK)
}