package ratelimit

import (
	"encoding/json"
	"os"
	"path/filepath"
	"time"

	"github.com/rs/zerolog/log"
)

const stateFile = ".anthropic-proxy-state.json"

type windowState struct {
	ResetsAt  time.Time `json:"resets_at"`
	TokensIn  int64     `json:"tokens_in"`
	TokensOut int64     `json:"tokens_out"`
}

type persistedState struct {
	FiveHour windowState `json:"five_hour"`
	SevenDay windowState `json:"seven_day"`
}

func statePath() string {
	home, err := os.UserHomeDir()
	if err != nil {
		return stateFile
	}
	return filepath.Join(home, ".claude", stateFile)
}

func (t *Tracker) Save() {
	t.mu.RLock()
	state := persistedState{
		FiveHour: windowState{
			ResetsAt:  t.fiveHour.ResetsAt,
			TokensIn:  t.fiveHour.TokensIn.Load(),
			TokensOut: t.fiveHour.TokensOut.Load(),
		},
		SevenDay: windowState{
			ResetsAt:  t.sevenDay.ResetsAt,
			TokensIn:  t.sevenDay.TokensIn.Load(),
			TokensOut: t.sevenDay.TokensOut.Load(),
		},
	}
	t.mu.RUnlock()

	data, err := json.Marshal(state)
	if err != nil {
		log.Warn().Err(err).Msg("failed to marshal tracker state")
		return
	}

	path := statePath()
	if err := os.WriteFile(path, data, 0644); err != nil {
		log.Warn().Err(err).Str("path", path).Msg("failed to save tracker state")
		return
	}
}

func (t *Tracker) Restore() {
	path := statePath()
	data, err := os.ReadFile(path)
	if err != nil {
		return
	}

	var state persistedState
	if err := json.Unmarshal(data, &state); err != nil {
		log.Warn().Err(err).Msg("failed to parse tracker state")
		return
	}

	t.mu.Lock()
	defer t.mu.Unlock()

	if !state.FiveHour.ResetsAt.IsZero() && state.FiveHour.ResetsAt == t.fiveHour.ResetsAt {
		t.fiveHour.TokensIn.Store(state.FiveHour.TokensIn)
		t.fiveHour.TokensOut.Store(state.FiveHour.TokensOut)
		log.Info().Int64("tokens_in", state.FiveHour.TokensIn).Int64("tokens_out", state.FiveHour.TokensOut).Msg("restored 5h window tokens")
	}

	if !state.SevenDay.ResetsAt.IsZero() && state.SevenDay.ResetsAt == t.sevenDay.ResetsAt {
		t.sevenDay.TokensIn.Store(state.SevenDay.TokensIn)
		t.sevenDay.TokensOut.Store(state.SevenDay.TokensOut)
		log.Info().Int64("tokens_in", state.SevenDay.TokensIn).Int64("tokens_out", state.SevenDay.TokensOut).Msg("restored 7d window tokens")
	}
}