package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"
	"time"

	"github.com/valyala/fasthttp"
)

const (
	baseURL           = "https://api.worldquantbrain.com"
	zeroStreakThreshold = 5 * 252
	requiredDays      = 2920
)

type Client struct {
	client   *fasthttp.Client
	username string
	password string
}

type AlphaRecord struct {
	ID           string  `json:"id"`
	Name         string  `json:"name"`
	DateCreated  string  `json:"dateCreated"`
	Sharpe       float64 `json:"sharpe"`
	Fitness      float64 `json:"fitness"`
	Turnover     float64 `json:"turnover"`
	Margin       float64 `json:"margin"`
	LongCount    float64 `json:"longCount"`
	ShortCount   float64 `json:"shortCount"`
	Decay        int     `json:"decay"`
	Code         string  `json:"code"`
}

type AlphaResponse struct {
	Results []struct {
		ID          string                 `json:"id"`
		Name        string                 `json:"name"`
		DateCreated string                 `json:"dateCreated"`
		Is          map[string]interface{} `json:"is"`
		Settings    struct {
			Decay int `json:"decay"`
		} `json:"settings"`
		Regular struct {
			Code string `json:"code"`
		} `json:"regular"`
	} `json:"results"`
}

type PnlResponse struct {
	Records [][]interface{} `json:"records"`
}

func NewClient(username, password string) *Client {
	return &Client{
		client:   &fasthttp.Client{},
		username: username,
		password: password,
	}
}

func (c *Client) getAuthHeader() string {
	auth := base64.StdEncoding.EncodeToString([]byte(c.username + ":" + c.password))
	return "Basic " + auth
}

func (c *Client) Login() error {
	req := fasthttp.AcquireRequest()
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseRequest(req)
	defer fasthttp.ReleaseResponse(resp)

	req.SetRequestURI(baseURL + "/authentication")
	req.Header.SetMethod("POST")
	req.Header.Set("Authorization", c.getAuthHeader())

	err := c.client.Do(req, resp)
	if err != nil {
		return err
	}

	fmt.Println(string(resp.Body()))
	return nil
}

func (c *Client) WaitGet(url string, maxRetries int) (*fasthttp.Response, error) {
	retries := 0
	for retries < maxRetries {
		for {
			req := fasthttp.AcquireRequest()
			resp := fasthttp.AcquireResponse()
			req.SetRequestURI(url)
			req.Header.SetMethod("GET")
			req.Header.Set("Authorization", c.getAuthHeader())

			err := c.client.Do(req, resp)
			if err != nil {
				fasthttp.ReleaseRequest(req)
				fasthttp.ReleaseResponse(resp)
				return nil, err
			}

			retryAfter := resp.Header.Peek("Retry-After")
			if len(retryAfter) == 0 {
				fasthttp.ReleaseRequest(req)
				if resp.StatusCode() < 400 {
					return resp, nil
				}
				fasthttp.ReleaseResponse(resp)
				break
			}

			sleepSec, _ := strconv.ParseFloat(string(retryAfter), 64)
			time.Sleep(time.Duration(sleepSec) * time.Second)
			fasthttp.ReleaseRequest(req)
			fasthttp.ReleaseResponse(resp)
		}

		time.Sleep(time.Duration(math.Pow(2, float64(retries))) * time.Second)
		retries++
	}

	return nil, fmt.Errorf("max retries exceeded")
}

func (c *Client) Get(url string) (*fasthttp.Response, error) {
	req := fasthttp.AcquireRequest()
	resp := fasthttp.AcquireResponse()
	req.SetRequestURI(url)
	req.Header.SetMethod("GET")
	req.Header.Set("Authorization", c.getAuthHeader())

	err := c.client.Do(req, resp)
	if err != nil {
		fasthttp.ReleaseRequest(req)
		fasthttp.ReleaseResponse(resp)
		return nil, err
	}

	return resp, nil
}

func (c *Client) Patch(url string, data map[string]interface{}) (*fasthttp.Response, error) {
	req := fasthttp.AcquireRequest()
	resp := fasthttp.AcquireResponse()
	defer fasthttp.ReleaseRequest(req)

	req.SetRequestURI(url)
	req.Header.SetMethod("PATCH")
	req.Header.Set("Authorization", c.getAuthHeader())
	req.Header.SetContentType("application/json")

	jsonData, _ := json.Marshal(data)
	req.SetBody(jsonData)

	err := c.client.Do(req, resp)
	if err != nil {
		fasthttp.ReleaseResponse(resp)
		return nil, err
	}

	return resp, nil
}

func GetAlphas(c *Client, startDate, endDate string, sharpeTh, fitnessTh float64, region string, alphaNum int, usage string) ([][]interface{}, *Client, error) {
	output := make([][]interface{}, 0)
	count := 0

	for i := 0; i < alphaNum; i += 100 {
		fmt.Println(i)

		urlE := fmt.Sprintf("%s/users/self/alphas?limit=100&offset=%d&status=UNSUBMITTED%%1FIS_FAIL&dateCreated%%3E=2025-%sT00:00:00-04:00&dateCreated%%3C2025-%sT00:00:00-04:00&is.fitness%%3E%f&is.sharpe%%3E%f&settings.region=%s&order=-is.sharpe&hidden=false&type!=SUPER",
			baseURL, i, startDate, endDate, fitnessTh, sharpeTh, region)

		urlC := fmt.Sprintf("%s/users/self/alphas?limit=100&offset=%d&status=UNSUBMITTED%%1FIS_FAIL&dateCreated%%3E=2025-%sT00:00:00-04:00&dateCreated%%3C2025-%sT00:00:00-04:00&is.fitness%%3C-%f&is.sharpe%%3C-%f&settings.region=%s&order=is.sharpe&hidden=false&type!=SUPER",
			baseURL, i, startDate, endDate, fitnessTh, sharpeTh, region)

		urls := []string{urlE}
		if usage != "submit" {
			urls = append(urls, urlC)
		}

		for _, url := range urls {
			resp, err := c.Get(url)
			if err != nil {
				fmt.Printf("%d finished re-login\n", i)
				c.Login()
				continue
			}

			var alphaResp AlphaResponse
			if err := json.Unmarshal(resp.Body(), &alphaResp); err != nil {
				fasthttp.ReleaseResponse(resp)
				fmt.Printf("%d finished re-login\n", i)
				c.Login()
				continue
			}
			fasthttp.ReleaseResponse(resp)

			for _, item := range alphaResp.Results {
				alphaID := item.ID
				name := item.Name
				dateCreated := item.DateCreated
				sharpe := getFloat(item.Is, "sharpe")
				fitness := getFloat(item.Is, "fitness")
				turnover := getFloat(item.Is, "turnover")
				margin := getFloat(item.Is, "margin")
				longCount := getFloat(item.Is, "longCount")
				shortCount := getFloat(item.Is, "shortCount")
				decay := item.Settings.Decay
				exp := item.Regular.Code

				count++

				if (longCount + shortCount) > 100 {
					if sharpe < -sharpeTh {
						exp = "-" + exp
					}

					rec := []interface{}{alphaID, exp, sharpe, turnover, fitness, margin, dateCreated, decay}
					fmt.Println(rec)

					if turnover > 0.7 {
						rec = append(rec, float64(decay)*4)
					} else if turnover > 0.6 {
						rec = append(rec, float64(decay)*3+3)
					} else if turnover > 0.5 {
						rec = append(rec, float64(decay)*3)
					} else if turnover > 0.4 {
						rec = append(rec, float64(decay)*2)
					} else if turnover > 0.35 {
						rec = append(rec, float64(decay)+4)
					} else if turnover > 0.3 {
						rec = append(rec, float64(decay)+2)
					}

					output = append(output, rec)
				}
			}
		}
	}

	fmt.Printf("count: %d\n", count)
	return output, c, nil
}

func getFloat(m map[string]interface{}, key string) float64 {
	if v, ok := m[key]; ok {
		switch val := v.(type) {
		case float64:
			return val
		case int:
			return float64(val)
		}
	}
	return 0
}

func CheckConsecutiveNonZeroValues(alphaID string, data [][]interface{}, requiredStreak int) bool {
	if len(data) < requiredStreak {
		return true
	}

	checkColumn := func(columnData []float64) bool {
		if len(columnData) < requiredStreak {
			return true
		}

		currentStreakCount := 0
		var currentStreakValue interface{}

		for _, value := range columnData {
			if value != 0 {
				if currentStreakValue != nil && value == currentStreakValue {
					currentStreakCount++
				} else {
					currentStreakValue = value
					currentStreakCount = 1
				}
			} else {
				currentStreakValue = nil
				currentStreakCount = 0
			}

			if currentStreakCount >= requiredStreak {
				return false
			}
		}
		return true
	}

	var column1Values, column2Values []float64
	for _, row := range data {
		if len(row) >= 3 {
			if v, ok := row[1].(float64); ok {
				column1Values = append(column1Values, v)
			}
			if v, ok := row[2].(float64); ok {
				column2Values = append(column2Values, v)
			}
		}
	}

	if len(column1Values) > 0 && len(column2Values) > 0 {
		isCol1AllZeros := allZeros(column1Values)
		isCol2AllZeros := allZeros(column2Values)
		if isCol1AllZeros || isCol2AllZeros {
			fmt.Println(alphaID, "不合法")
			return false
		}
	}

	if !checkColumn(column1Values) {
		fmt.Println(alphaID, "不合法")
		return false
	}

	if !checkColumn(column2Values) {
		fmt.Println(alphaID, "不合法")
		return false
	}

	return true
}

func allZeros(arr []float64) bool {
	for _, v := range arr {
		if v != 0 {
			return false
		}
	}
	return true
}

func GetAlphaPnlLegal(c *Client, alphaID string) bool {
	notLegalID := make([]string, 0)

	url := baseURL + "/alphas/" + alphaID + "/recordsets/pnl"
	resp, err := c.WaitGet(url, 10)
	if err != nil {
		return false
	}
	defer fasthttp.ReleaseResponse(resp)

	var pnlResp PnlResponse
	if err := json.Unmarshal(resp.Body(), &pnlResp); err != nil {
		return false
	}

	records := pnlResp.Records
	if len(records) == 0 {
		return false
	}

	var dateList []time.Time
	for _, record := range records {
		if len(record) == 0 {
			continue
		}
		dateStr, ok := record[0].(string)
		if !ok {
			return false
		}
		dateObj, err := time.Parse("2006-01-02", dateStr)
		if err != nil {
			return false
		}
		dateList = append(dateList, dateObj)
	}

	if len(dateList) == 0 {
		return false
	}

	minDate := dateList[0]
	maxDate := dateList[0]
	for _, d := range dateList {
		if d.Before(minDate) {
			minDate = d
		}
		if d.After(maxDate) {
			maxDate = d
		}
	}

	totalDays := int(maxDate.Sub(minDate).Hours() / 24)
	if totalDays < requiredDays {
		return false
	}

	col1Zeros := make([]bool, 0)
	for _, record := range records {
		if len(record) >= 2 {
			if v, ok := record[1].(float64); ok {
				col1Zeros = append(col1Zeros, v == 0)
			}
		}
	}

	col1MaxZeroStreak := maxConsecutiveZeros(col1Zeros)

	if col1MaxZeroStreak >= zeroStreakThreshold {
		fmt.Printf("%s 不合法：存在连续%d年零值\n", alphaID, zeroStreakThreshold/252)
		notLegalID = append(notLegalID, alphaID)
		return false
	}

	if !CheckConsecutiveNonZeroValues(alphaID, records, 200) {
		return false
	}

	_ = notLegalID
	return true
}

func maxConsecutiveZeros(arr []bool) int {
	maxStreak := 0
	currentStreak := 0
	for _, val := range arr {
		if val {
			currentStreak++
			if currentStreak > maxStreak {
				maxStreak = currentStreak
			}
		} else {
			currentStreak = 0
		}
	}
	return maxStreak
}

func Mute(c *Client, alphaID string) {
	url := baseURL + "/alphas/" + alphaID
	data := map[string]interface{}{
		"hidden": true,
	}
	c.Patch(url, data)
}

func main() {
	client := NewClient("", "")
	client.Login()

	foTracker, c, err := GetAlphas(client, "12-01", "12-31", 1, 0.5, "USA", 1000, "submit")
	if err != nil {
		fmt.Println("Error:", err)
		return
	}

	fNum := len(foTracker)
	fmt.Printf("%d 个alpha 进行pnl合法检测，请耐心等待\n", fNum)
	fmt.Println(len(foTracker))

	count := 0
	for i := len(foTracker) - 1; i >= 0; i-- {
		if count%25 == 0 {
			fmt.Printf("=========== %d ===========\n", count)
		}
		count++

		alphaID, ok := foTracker[i][0].(string)
		if !ok {
			continue
		}

		if !GetAlphaPnlLegal(c, alphaID) {
			fmt.Println(alphaID, "已经隐藏")
			Mute(c, alphaID)
		}
	}
}