package state

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"sync"
	"time"
)

const PaletteSize = 16

type Row []uint8

func (r Row) MarshalJSON() ([]byte, error) {
	vals := make([]int, len(r))
	for i, v := range r {
		vals[i] = int(v)
	}
	return json.Marshal(vals)
}

func (r *Row) UnmarshalJSON(data []byte) error {
	var vals []int
	if err := json.Unmarshal(data, &vals); err != nil {
		return err
	}
	row := make(Row, len(vals))
	for i, v := range vals {
		if v < 0 || v > 255 {
			return fmt.Errorf("pixel value %d out of range", v)
		}
		row[i] = uint8(v)
	}
	*r = row
	return nil
}

var (
	ErrOutOfBounds  = errors.New("coordinates out of bounds")
	ErrInvalidColor = errors.New("invalid color index")
	ErrImageTooBig  = errors.New("image too large")
	ErrEmptyMessage = errors.New("message is empty")
	ErrMessageLong  = errors.New("message too long")
)

type Message struct {
	Text string    `json:"text"`
	At   time.Time `json:"at"`
}

type State struct {
	mu          sync.RWMutex
	Width       int       `json:"width"`
	Height      int       `json:"height"`
	Pixels      []Row `json:"pixels"`
	Messages    []Message `json:"messages"`
	maxMessages int
	maxMsgLen   int
	maxImageDim int
}

func New(width, height int) *State {
	pixels := make([]Row, height)
	for i := range pixels {
		pixels[i] = make(Row, width)
	}
	return &State{
		Width:       width,
		Height:      height,
		Pixels:      pixels,
		Messages:    []Message{},
		maxMessages: 50,
		maxMsgLen:   140,
		maxImageDim: 32,
	}
}

func (s *State) CheckPixel(x, y int, color uint8) error {
	if x < 0 || x >= s.Width || y < 0 || y >= s.Height {
		return ErrOutOfBounds
	}
	if color >= PaletteSize {
		return ErrInvalidColor
	}
	return nil
}

func (s *State) SetPixel(x, y int, color uint8) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := s.CheckPixel(x, y, color); err != nil {
		return err
	}
	s.Pixels[y][x] = color
	return nil
}

func (s *State) StampImage(x, y int, pixels []Row) error {
	if len(pixels) == 0 || len(pixels) > s.maxImageDim {
		return ErrImageTooBig
	}
	for _, row := range pixels {
		if len(row) == 0 || len(row) > s.maxImageDim {
			return ErrImageTooBig
		}
		for _, c := range row {
			if c >= PaletteSize {
				return ErrInvalidColor
			}
		}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	for dy, row := range pixels {
		for dx, c := range row {
			px, py := x+dx, y+dy
			if px >= 0 && px < s.Width && py >= 0 && py < s.Height {
				s.Pixels[py][px] = c
			}
		}
	}
	return nil
}

func (s *State) AddMessage(text string) error {
	if text == "" {
		return ErrEmptyMessage
	}
	if len(text) > s.maxMsgLen {
		return ErrMessageLong
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Messages = append(s.Messages, Message{Text: text, At: time.Now().UTC()})
	if len(s.Messages) > s.maxMessages {
		s.Messages = s.Messages[len(s.Messages)-s.maxMessages:]
	}
	return nil
}

type Snapshot struct {
	Width    int       `json:"width"`
	Height   int       `json:"height"`
	Pixels   []Row `json:"pixels"`
	Messages []Message `json:"messages"`
}

func (s *State) Snapshot() Snapshot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	pixels := make([]Row, s.Height)
	for i, row := range s.Pixels {
		cp := make(Row, len(row))
		copy(cp, row)
		pixels[i] = cp
	}
	msgs := make([]Message, len(s.Messages))
	copy(msgs, s.Messages)
	return Snapshot{Width: s.Width, Height: s.Height, Pixels: pixels, Messages: msgs}
}

func (s *State) Save(path string) error {
	snap := s.Snapshot()
	data, err := json.Marshal(snap)
	if err != nil {
		return err
	}
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

func Load(path string) (*State, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var snap Snapshot
	if err := json.Unmarshal(data, &snap); err != nil {
		return nil, err
	}
	if snap.Width <= 0 || snap.Height <= 0 || len(snap.Pixels) != snap.Height {
		return nil, fmt.Errorf("invalid snapshot dimensions")
	}
	for _, row := range snap.Pixels {
		if len(row) != snap.Width {
			return nil, fmt.Errorf("invalid snapshot row width")
		}
	}
	st := New(snap.Width, snap.Height)
	st.Pixels = snap.Pixels
	if snap.Messages != nil {
		st.Messages = snap.Messages
	}
	return st, nil
}