package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

type TrafficEvent struct {
	Monitoring bool   `json:"monitoring"`
	ID         string `json:"id"`
	Sequence   int64  `json:"sequence,omitempty"`
	Time       int64  `json:"time"`
	Kind       string `json:"kind"`
	IP         string `json:"ip"`
	Domain     string `json:"domain,omitempty"`
	Method     string `json:"method,omitempty"`
	Path       string `json:"path,omitempty"`
	Status     int    `json:"status,omitempty"`
	Reason     string `json:"reason,omitempty"`
	Country    string `json:"country,omitempty"`
	Agent      string `json:"agent,omitempty"`
}

type TrafficBatch struct {
	Events          []TrafficEvent `json:"events"`
	Files           int            `json:"files"`
	BacklogBytes    int64          `json:"backlogBytes"`
	CrowdSecLimited bool           `json:"crowdsecLimited"`
	ParseErrors     uint64         `json:"parseErrors"`
	LogError        string         `json:"logError,omitempty"`
	CrowdSecError   string         `json:"crowdsecError,omitempty"`
	CrowdSecAt      int64          `json:"crowdsecAt"`
	StartedAt       int64          `json:"startedAt"`
}

type TrafficSource struct {
	Agent    string       `json:"agent"`
	Name     string       `json:"name"`
	LastSeen int64        `json:"lastSeen"`
	Status   TrafficBatch `json:"status"`
}

type TrafficCount struct {
	Name  string `json:"name"`
	Count int64  `json:"count"`
}

type TrafficCountryDetail struct {
	Country   string `json:"country"`
	Requests  int64  `json:"requests"`
	UniqueIPs int64  `json:"uniqueIps"`
	Alerts    int64  `json:"alerts"`
	Decisions int64  `json:"decisions"`
	LastSeen  int64  `json:"lastSeen"`
}
type TrafficView struct {
	CountryDetails []TrafficCountryDetail `json:"countryDetails"`

	Network   *NetworkMapView `json:"network,omitempty"`
	Events    []TrafficEvent  `json:"events"`
	Sources   []TrafficSource `json:"sources"`
	Countries []TrafficCount  `json:"countries"`
	Domains   []TrafficCount  `json:"domains"`
	IPs       []TrafficCount  `json:"ips"`
	Requests  int64           `json:"requests"`
	UniqueIPs int64           `json:"uniqueIps"`
	Alerts    int64           `json:"alerts"`
	Decisions int64           `json:"decisions"`
	QPS       float64         `json:"qps"`
	Series    []TrafficCount  `json:"series"`
	GeoReady  bool            `json:"geoReady"`
	GeoDate   string          `json:"geoDate"`
	Now       int64           `json:"now"`
	Settings  TrafficSettings `json:"settings"`
	Limited   bool            `json:"limited"`
}

const trafficSchema = `CREATE TABLE IF NOT EXISTS traffic_events (
 seq INTEGER PRIMARY KEY AUTOINCREMENT, agent TEXT NOT NULL, event_id TEXT NOT NULL,
 ts INTEGER NOT NULL, kind TEXT NOT NULL, ip TEXT NOT NULL, domain TEXT NOT NULL,
 method TEXT NOT NULL, path TEXT NOT NULL, status INTEGER NOT NULL, reason TEXT NOT NULL,
 country TEXT NOT NULL, UNIQUE(agent,event_id));
 CREATE INDEX IF NOT EXISTS traffic_ts ON traffic_events(ts);
 CREATE INDEX IF NOT EXISTS traffic_filter ON traffic_events(domain,kind,ts);
 CREATE TABLE IF NOT EXISTS network_sources(agent TEXT PRIMARY KEY, snapshot TEXT NOT NULL);
 CREATE TABLE IF NOT EXISTS traffic_sources(agent TEXT PRIMARY KEY, last_seen INTEGER NOT NULL, status_json TEXT NOT NULL);`

func (s *Store) SaveTraffic(agent string, batch TrafficBatch) error {
	if len(batch.Events) > 500 {
		return errors.New("too many traffic events")
	}
	now := time.Now().Unix()
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	stmt, err := tx.Prepare(`INSERT OR IGNORE INTO traffic_events(agent,event_id,ts,kind,ip,domain,method,path,status,reason,country) VALUES(?,?,?,?,?,?,?,?,?,?,?)`)
	if err != nil {
		return err
	}
	defer stmt.Close()
	for _, e := range batch.Events {
		if e.Kind != "request" && e.Kind != "alert" && e.Kind != "decision" {
			return errors.New("invalid traffic kind")
		}
		if net.ParseIP(e.IP) == nil || len(e.ID) == 0 || len(e.ID) > 160 {
			return errors.New("invalid traffic event")
		}
		if e.Time < now-86400 || e.Time > now+300 {
			continue
		}
		e.IP = net.ParseIP(e.IP).String()
		country := lookupTrafficCountry(e.IP)
		_, err = stmt.Exec(agent, e.ID, e.Time, e.Kind, e.IP, trafficText(e.Domain, 253), trafficText(e.Method, 16), trafficPath(e.Path), e.Status, trafficText(e.Reason, 256), country)
		if err != nil {
			return err
		}
	}
	batch.Events = nil
	raw, err := json.Marshal(batch)
	if err != nil {
		return err
	}
	if _, err = tx.Exec(`INSERT INTO traffic_sources(agent,last_seen,status_json) VALUES(?,?,?) ON CONFLICT(agent) DO UPDATE SET last_seen=excluded.last_seen,status_json=excluded.status_json`, agent, now, string(raw)); err != nil {
		return err
	}
	// Bound both age and total rows, including a high-traffic installation.
	if _, err = tx.Exec(`DELETE FROM traffic_events WHERE ts<? OR seq <= COALESCE((SELECT seq FROM traffic_events ORDER BY seq DESC LIMIT 1 OFFSET 200000),0)`, now-86400); err != nil {
		return err
	}
	return tx.Commit()
}

func trafficText(s string, max int) string {
	s = strings.Map(func(r rune) rune {
		if r < 32 || r == 127 {
			return -1
		}
		return r
	}, s)
	r := []rune(s)
	if len(r) > max {
		return string(r[:max])
	}
	return s
}
func trafficPath(s string) string {
	s = strings.SplitN(s, "?", 2)[0]
	s = strings.SplitN(s, "#", 2)[0]
	return trafficText(s, 256)
}

func (s *Store) TrafficView(domain, kind string, external bool, hideMonitoring ...bool) (TrafficView, error) {
	now := time.Now().Unix()
	v := TrafficView{Events: []TrafficEvent{}, Sources: []TrafficSource{}, Countries: []TrafficCount{}, Domains: []TrafficCount{}, IPs: []TrafficCount{}, Series: []TrafficCount{}, Now: now, GeoReady: trafficGeoReady(), GeoDate: trafficGeoDate}
	v.Settings = s.TrafficSettings()
	hide := len(hideMonitoring) > 0 && hideMonitoring[0]
	if v.Settings.NetworkEnabled {
		snapshot, err := s.ReadNetworkSnapshot(v.Settings.AgentID, v.Settings.MonitoringRules, hide)
		if err != nil {
			snapshot.Error = "Не удалось прочитать сетевой снимок"
		}
		network := makeNetworkMap(snapshot, external)
		v.Network = &network
	}
	where := ` WHERE ts>=? AND agent=?`
	args := []any{now - 86400, v.Settings.AgentID}
	if hide {
		clause, params := monitoringSQL(v.Settings.MonitoringRules)
		where += clause
		args = append(args, params...)
	}
	if external {
		where += ` AND country!='LAN'`
	}
	if domain != "" {
		where += ` AND domain=?`
		args = append(args, domain)
	}
	if kind != "" {
		where += ` AND kind=?`
		args = append(args, kind)
	}
	err := s.db.QueryRow(`SELECT COALESCE(SUM(kind='request'),0),COUNT(DISTINCT CASE WHEN kind='request' THEN ip END),COALESCE(SUM(kind='alert'),0),COALESCE(SUM(kind='decision'),0),COALESCE(SUM(kind='request' AND ts>=?),0)/60.0 FROM traffic_events`+where, append([]any{now - 60}, args...)...).Scan(&v.Requests, &v.UniqueIPs, &v.Alerts, &v.Decisions, &v.QPS)
	if err != nil {
		return v, err
	}
	rows, err := s.db.Query(`SELECT seq,event_id,ts,kind,ip,domain,method,path,status,reason,country,agent FROM traffic_events`+where+` ORDER BY ts DESC,seq DESC LIMIT 100`, args...)
	if err != nil {
		return v, err
	}
	for rows.Next() {
		var e TrafficEvent
		if err = rows.Scan(&e.Sequence, &e.ID, &e.Time, &e.Kind, &e.IP, &e.Domain, &e.Method, &e.Path, &e.Status, &e.Reason, &e.Country, &e.Agent); err != nil {
			rows.Close()
			return v, err
		}
		e.Monitoring = isMonitoringRequest(e, v.Settings.MonitoringRules)
		v.Events = append(v.Events, e)
	}
	err = rows.Err()
	rows.Close()
	if err != nil {
		return v, err
	}
	for _, q := range []struct {
		column string
		out    *[]TrafficCount
	}{{"country", &v.Countries}, {"domain", &v.Domains}, {"ip", &v.IPs}} {
		r, e := s.db.Query(`SELECT `+q.column+`,COUNT(*) FROM traffic_events`+where+` GROUP BY `+q.column+` ORDER BY COUNT(*) DESC LIMIT 50`, args...)
		if e != nil {
			return v, e
		}
		for r.Next() {
			var c TrafficCount
			if e = r.Scan(&c.Name, &c.Count); e != nil {
				r.Close()
				return v, e
			}
			*q.out = append(*q.out, c)
		}
		e = r.Err()
		r.Close()
		if e != nil {
			return v, e
		}
	}
	detailRows, detailErr := s.db.Query(`SELECT country,SUM(kind='request'),COUNT(DISTINCT CASE WHEN kind='request' THEN ip END),SUM(kind='alert'),SUM(kind='decision'),MAX(ts) FROM traffic_events`+where+` GROUP BY country`, args...)
	if detailErr != nil {
		return v, detailErr
	}
	v.CountryDetails = []TrafficCountryDetail{}
	for detailRows.Next() {
		var d TrafficCountryDetail
		if err = detailRows.Scan(&d.Country, &d.Requests, &d.UniqueIPs, &d.Alerts, &d.Decisions, &d.LastSeen); err != nil {
			detailRows.Close()
			return v, err
		}
		v.CountryDetails = append(v.CountryDetails, d)
	}
	detailErr = detailRows.Err()
	detailRows.Close()
	if detailErr != nil {
		return v, detailErr
	}
	r, err := s.db.Query(`SELECT CAST((ts/60)*60 AS TEXT),COUNT(*) FROM traffic_events`+where+` AND kind='request' AND ts>=? GROUP BY ts/60 ORDER BY ts/60`, append(args, now-3600)...)
	if err != nil {
		return v, err
	}
	for r.Next() {
		var c TrafficCount
		if err = r.Scan(&c.Name, &c.Count); err != nil {
			r.Close()
			return v, err
		}
		v.Series = append(v.Series, c)
	}
	err = r.Err()
	r.Close()
	if err != nil {
		return v, err
	}
	r, err = s.db.Query(`SELECT t.agent,COALESCE(a.name,t.agent),t.last_seen,t.status_json FROM traffic_sources t LEFT JOIN agents a ON a.id=t.agent WHERE t.agent=? ORDER BY t.last_seen DESC`, v.Settings.AgentID)
	if err != nil {
		return v, err
	}
	for r.Next() {
		var src TrafficSource
		var raw string
		if err = r.Scan(&src.Agent, &src.Name, &src.LastSeen, &raw); err != nil {
			r.Close()
			return v, err
		}
		json.Unmarshal([]byte(raw), &src.Status)
		v.Sources = append(v.Sources, src)
	}
	err = r.Err()
	r.Close()
	var count int64
	s.db.QueryRow(`SELECT COUNT(*) FROM traffic_events`).Scan(&count)
	v.Limited = count >= 200000
	return v, err
}

func trafficRoutes(mux *http.ServeMux, s *Store) {
	trafficSettingsRoute(mux, s)
	networkRoutes(mux, s)
	mux.HandleFunc("/traffic", func(w http.ResponseWriter, r *http.Request) {
		if !s.TrafficSettings().Enabled {
			http.Error(w, "Карта трафика отключена. Включите её в настройках Dashboard.", 404)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write([]byte(trafficHTML))
	})
	mux.HandleFunc("/api/traffic", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			w.WriteHeader(405)
			return
		}
		if !s.TrafficSettings().Enabled {
			http.Error(w, "Карта отключена", 404)
			return
		}
		view, err := s.TrafficView(r.URL.Query().Get("domain"), r.URL.Query().Get("kind"), r.URL.Query().Get("external") == "1", r.URL.Query().Get("hideMonitoring") == "1")
		writeJSON(w, view, err)
	})
	mux.HandleFunc("/api/traffic/world", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("Cache-Control", "public, max-age=86400")
		w.Write(trafficWorld)
	})
	mux.HandleFunc("/api/agent/traffic", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "POST" {
			w.WriteHeader(405)
			return
		}
		id, secret, ok := agentCredentials(r)
		if !ok || !s.AuthenticateAgent(id, secret) {
			http.Error(w, "Unauthorized", 401)
			return
		}
		settings := s.TrafficSettings()
		if !settings.Enabled || settings.AgentID != id {
			http.Error(w, "Источник отключён", 403)
			return
		}
		var batch TrafficBatch
		if json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&batch) != nil {
			http.Error(w, "Invalid traffic batch", 400)
			return
		}
		if err := s.SaveTraffic(id, batch); err != nil {
			http.Error(w, fmt.Sprint(err), 500)
			return
		}
		writeJSON(w, map[string]bool{"accepted": true}, nil)
	})
}