package collector

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"time"
)

type LoginEvent struct {
	Timestamp string `json:"timestamp"`
	User      string `json:"user"`
	Source    string `json:"source"`
	Port      string `json:"port"`
	Method    string `json:"method"`
	Status    string `json:"status"`
	Key       string `json:"key,omitempty"`
}

type Session struct {
	User     string `json:"user"`
	TTY      string `json:"tty"`
	From     string `json:"from"`
	Login    string `json:"login"`
	Logout   string `json:"logout"`
	Duration string `json:"duration"`
}

type SecuritySummary struct {
	TotalAccepted int    `json:"total_accepted"`
	TotalFailed   int    `json:"total_failed"`
	UniqueIPs     int    `json:"unique_ips"`
	UniqueUsers   int    `json:"unique_users"`
	LastLogin     string `json:"last_login"`
}

type SecurityData struct {
	SSHLogins    []LoginEvent    `json:"ssh_logins"`
	WebGUILogins []LoginEvent    `json:"webgui_logins"`
	Sessions     []Session       `json:"sessions"`
	Summary      SecuritySummary `json:"summary"`
}

var (
	reAccepted = regexp.MustCompile(`Accepted\s+(publickey|password)\s+for\s+(\S+)\s+from\s+(\S+)\s+port\s+(\d+)`)
	reFailed   = regexp.MustCompile(`Failed\s+password\s+for\s+(?:invalid user\s+)?(\S+)\s+from\s+(\S+)\s+port\s+(\d+)`)
	reInvalid  = regexp.MustCompile(`Invalid user\s+(\S+)\s+from\s+(\S+)\s+port\s+(\d+)`)
	// RFC 5424 Syslog: <NN>1 2024-01-15T10:30:00+01:00 ... (with version digit)
	reTimestamp = regexp.MustCompile(`^<\d+>(?:1\s+)?(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}[^\s]*)`)
	// BSD syslog: Jan 15 10:30:00
	reBSDTimestamp = regexp.MustCompile(`^<\d+>\s*([A-Z][a-z]{2}\s+\d+\s+\d{2}:\d{2}:\d{2})`)
)

func CollectSecurity() *SecurityData {
	data := &SecurityData{
		SSHLogins:    []LoginEvent{},
		WebGUILogins: []LoginEvent{},
		Sessions:     []Session{},
	}

	// Collect audit log files
	logFiles := []string{"/var/log/audit/latest.log"}
	now := time.Now()
	for i := 0; i < 7; i++ {
		d := now.AddDate(0, 0, -i)
		logFiles = append(logFiles, fmt.Sprintf("/var/log/audit/audit_%s.log", d.Format("20060102")))
	}

	seen := make(map[string]bool)
	for _, path := range logFiles {
		events, err := parseAuditLog(path)
		if err != nil {
			log.Printf("Security: %s: %v", path, err)
			continue
		}
		log.Printf("Security: %s: %d events", path, len(events))
		for _, e := range events {
			key := e.Timestamp + e.User + e.Source + e.Port + e.Status
			if !seen[key] {
				seen[key] = true
				data.SSHLogins = append(data.SSHLogins, e)
			}
		}
	}

	// Sort by timestamp descending, limit to 100
	sort.Slice(data.SSHLogins, func(i, j int) bool {
		return data.SSHLogins[i].Timestamp > data.SSHLogins[j].Timestamp
	})
	if len(data.SSHLogins) > 100 {
		data.SSHLogins = data.SSHLogins[:100]
	}

	// Collect sessions
	data.Sessions = collectSessions()

	// Build summary
	ips := make(map[string]bool)
	users := make(map[string]bool)
	for _, e := range data.SSHLogins {
		if e.Status == "accepted" {
			data.Summary.TotalAccepted++
		} else {
			data.Summary.TotalFailed++
		}
		if e.Source != "" {
			ips[e.Source] = true
		}
		if e.User != "" {
			users[e.User] = true
		}
	}
	data.Summary.UniqueIPs = len(ips)
	data.Summary.UniqueUsers = len(users)
	if len(data.SSHLogins) > 0 {
		data.Summary.LastLogin = data.SSHLogins[0].Timestamp
	}

	return data
}

func parseAuditLog(path string) ([]LoginEvent, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var events []LoginEvent
	scanner := bufio.NewScanner(f)
	// Increase buffer for long lines
	scanner.Buffer(make([]byte, 0, 64*1024), 256*1024)

	for scanner.Scan() {
		line := scanner.Text()

		// Extract timestamp
		ts := extractTimestamp(line)

		if m := reAccepted.FindStringSubmatch(line); m != nil {
			events = append(events, LoginEvent{
				Timestamp: ts,
				User:      m[2],
				Source:    m[3],
				Port:      m[4],
				Method:    m[1],
				Status:    "accepted",
			})
		} else if m := reFailed.FindStringSubmatch(line); m != nil {
			events = append(events, LoginEvent{
				Timestamp: ts,
				User:      m[1],
				Source:    m[2],
				Port:      m[3],
				Method:    "password",
				Status:    "failed",
			})
		} else if m := reInvalid.FindStringSubmatch(line); m != nil {
			events = append(events, LoginEvent{
				Timestamp: ts,
				User:      m[1],
				Source:    m[2],
				Port:      m[3],
				Method:    "invalid_user",
				Status:    "failed",
			})
		}
	}

	return events, nil
}

func extractTimestamp(line string) string {
	if m := reTimestamp.FindStringSubmatch(line); m != nil {
		return m[1]
	}
	if m := reBSDTimestamp.FindStringSubmatch(line); m != nil {
		// Parse BSD timestamp and add current year
		t, err := time.Parse("Jan 2 15:04:05", m[1])
		if err == nil {
			t = t.AddDate(time.Now().Year(), 0, 0)
			return t.Format(time.RFC3339)
		}
		return m[1]
	}
	return ""
}

func collectSessions() []Session {
	out, err := exec.Command("/usr/bin/last", "-100").Output()
	if err != nil {
		return nil
	}

	var sessions []Session
	for _, line := range strings.Split(string(out), "\n") {
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "wtmp") || strings.HasPrefix(line, "utx.log") {
			continue
		}

		fields := strings.Fields(line)
		if len(fields) < 4 {
			continue
		}

		s := Session{
			User: fields[0],
			TTY:  fields[1],
		}

		// Determine if there's a "from" field (IP or hostname)
		idx := 2
		if len(fields) > 4 && !isMonth(fields[2]) {
			s.From = fields[2]
			idx = 3
		}

		// Rest is login time, logout, duration
		rest := strings.Join(fields[idx:], " ")

		// Try to extract login/logout/duration from the rest
		if dashIdx := strings.Index(rest, " - "); dashIdx >= 0 {
			s.Login = strings.TrimSpace(rest[:dashIdx])
			after := rest[dashIdx+3:]
			if parenIdx := strings.Index(after, "("); parenIdx >= 0 {
				s.Logout = strings.TrimSpace(after[:parenIdx])
				dur := after[parenIdx:]
				dur = strings.Trim(dur, "()")
				s.Duration = dur
			} else {
				s.Logout = strings.TrimSpace(after)
			}
		} else {
			s.Login = rest
			if strings.Contains(rest, "still logged in") {
				s.Logout = "still logged in"
			}
		}

		sessions = append(sessions, s)
	}

	return sessions
}

func isMonth(s string) bool {
	months := []string{"Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"}
	for _, m := range months {
		if s == m {
			return true
		}
	}
	return false
}