package server

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
)

type Session struct {
	Username    string    `json:"username"`
	DisplayName string    `json:"display_name"`
	Role        string    `json:"role"`
	CreatedAt   time.Time `json:"created_at"`
	ExpiresAt   time.Time `json:"expires_at"`
	SessionID   string    `json:"session_id"`
}

type SessionManager struct {
	sessions map[string]Session
	mu       sync.RWMutex
	secret   []byte
	ttl      time.Duration
}

type sessionPayload struct {
	SessionID   string `json:"session_id"`
	Username    string `json:"username"`
	DisplayName string `json:"display_name"`
	Role        string `json:"role"`
	Exp         int64  `json:"exp"`
}

func NewSessionManager(secret []byte, ttl time.Duration) *SessionManager {
	return &SessionManager{
		sessions: make(map[string]Session),
		secret:   secret,
		ttl:      ttl,
	}
}

func (m *SessionManager) CreateSession(username, displayName, role string) (string, error) {
	sessionID := uuid.New().String()
	now := time.Now()
	expiresAt := now.Add(m.ttl)

	m.mu.Lock()
	m.sessions[sessionID] = Session{
		Username:    username,
		DisplayName: displayName,
		Role:        role,
		CreatedAt:   now,
		ExpiresAt:   expiresAt,
		SessionID:   sessionID,
	}
	m.mu.Unlock()

	return m.createSignedToken(sessionID, username, displayName, role, expiresAt.Unix())
}

func (m *SessionManager) createSignedToken(sessionID, username, displayName, role string, exp int64) (string, error) {
	payload := sessionPayload{
		SessionID:   sessionID,
		Username:    username,
		DisplayName: displayName,
		Role:        role,
		Exp:         exp,
	}

	payloadJSON, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}

	payloadB64 := base64.RawURLEncoding.EncodeToString(payloadJSON)
	
	h := hmac.New(sha256.New, m.secret)
	h.Write(payloadJSON)
	signature := h.Sum(nil)
	signatureB64 := base64.RawURLEncoding.EncodeToString(signature)

	return fmt.Sprintf("%s.%s", payloadB64, signatureB64), nil
}

func (m *SessionManager) ValidateSession(token string) (*Session, string, error) {
	parts := strings.Split(token, ".")
	if len(parts) != 2 {
		return nil, "", fmt.Errorf("invalid token format")
	}

	payloadB64 := parts[0]
	signatureB64 := parts[1]

	payloadJSON, err := base64.RawURLEncoding.DecodeString(payloadB64)
	if err != nil {
		return nil, "", err
	}

	signature, err := base64.RawURLEncoding.DecodeString(signatureB64)
	if err != nil {
		return nil, "", err
	}

	h := hmac.New(sha256.New, m.secret)
	h.Write(payloadJSON)
	if !hmac.Equal(signature, h.Sum(nil)) {
		return nil, "", fmt.Errorf("invalid signature")
	}

	var payload sessionPayload
	if err := json.Unmarshal(payloadJSON, &payload); err != nil {
		return nil, "", err
	}

	if time.Now().Unix() > payload.Exp {
		return nil, "", fmt.Errorf("token expired")
	}

	m.mu.RLock()
	session, ok := m.sessions[payload.SessionID]
	m.mu.RUnlock()

	if !ok {
		return nil, "", fmt.Errorf("session not found")
	}

	return &session, "", nil
}

func (m *SessionManager) RevokeSession(token string) {
	parts := strings.Split(token, ".")
	if len(parts) != 2 {
		return
	}

	payloadJSON, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return
	}

	var payload sessionPayload
	if err := json.Unmarshal(payloadJSON, &payload); err != nil {
		return
	}

	m.mu.Lock()
	delete(m.sessions, payload.SessionID)
	m.mu.Unlock()
}