package postgres

import (
	"context"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var ErrUnauthorized = errors.New("unauthorized")

type Store struct{ Pool *pgxpool.Pool }

type Agent struct {
	ID        string    `json:"id"`
	Hostname  string    `json:"hostname"`
	LastIP    string    `json:"last_ip"`
	Enabled   bool      `json:"enabled"`
	FirstSeen time.Time `json:"first_seen"`
	LastSeen  time.Time `json:"last_seen"`
}
type Detection struct {
	ID          int64     `json:"id"`
	Fingerprint string    `json:"fingerprint"`
	RuleName    string    `json:"rule_name"`
	Severity    string    `json:"severity"`
	Status      string    `json:"status"`
	Hostname    string    `json:"hostname"`
	UserName    string    `json:"user_name"`
	SourceIP    string    `json:"source_ip"`
	Workstation string    `json:"workstation"`
	Summary     string    `json:"summary"`
	EventCode   uint32    `json:"event_code"`
	Score       float64   `json:"score"`
	WindowStart time.Time `json:"window_start"`
	WindowEnd   time.Time `json:"window_end"`
	FirstSeen   time.Time `json:"first_seen"`
	LastSeen    time.Time `json:"last_seen"`
	Count       int64     `json:"count"`
}

func Open(ctx context.Context, url string) (*Store, error) {
	p, e := pgxpool.New(ctx, url)
	if e != nil {
		return nil, e
	}
	cctx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()
	if e = p.Ping(cctx); e != nil {
		p.Close()
		return nil, e
	}
	return &Store{Pool: p}, nil
}
func (s *Store) Close() { s.Pool.Close() }

func (s *Store) AuthenticateOrEnroll(ctx context.Context, tenant, hostname, apiKey, enrollmentKey, expectedEnrollment, remoteIP string) (string, error) {
	hostname = strings.TrimSpace(hostname)
	apiKey = strings.TrimSpace(apiKey)
	var id, hash string
	var enabled bool
	err := s.Pool.QueryRow(ctx, `SELECT id::text, api_key_hash, enabled FROM agents WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname).Scan(&id, &hash, &enabled)
	if err != nil {
		if !errors.Is(err, pgx.ErrNoRows) {
			return "", err
		}
		if enrollmentKey == "" || !secureEqual(hashHex(enrollmentKey), hashHex(expectedEnrollment)) {
			return "", ErrUnauthorized
		}
		err = s.Pool.QueryRow(ctx, `INSERT INTO agents(tenant_id,hostname,api_key_hash,last_ip) VALUES($1,$2,$3,$4) ON CONFLICT(tenant_id,hostname) DO NOTHING RETURNING id::text`, tenant, hostname, hashHex(apiKey), remoteIP).Scan(&id)
		if err == nil {
			return id, nil
		}
		if !errors.Is(err, pgx.ErrNoRows) {
			return "", err
		}
		// A concurrent first request may have enrolled the host between SELECT and INSERT.
		err = s.Pool.QueryRow(ctx, `SELECT id::text, api_key_hash, enabled FROM agents WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname).Scan(&id, &hash, &enabled)
		if err != nil || !enabled || !secureEqual(strings.ToLower(hash), hashHex(apiKey)) {
			return "", ErrUnauthorized
		}
		return id, nil
	}
	if !enabled || !secureEqual(strings.ToLower(hash), hashHex(apiKey)) {
		return "", ErrUnauthorized
	}
	_, err = s.Pool.Exec(ctx, `UPDATE agents SET last_seen=now(), last_ip=$3 WHERE tenant_id=$1 AND hostname=$2`, tenant, hostname, remoteIP)
	return id, err
}

func (s *Store) ListAgents(ctx context.Context, tenant string) ([]Agent, error) {
	rows, e := s.Pool.Query(ctx, `SELECT id::text,hostname,enabled,COALESCE(last_ip,''),first_seen,last_seen FROM agents WHERE tenant_id=$1 ORDER BY last_seen DESC`, tenant)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	var out []Agent
	for rows.Next() {
		var a Agent
		if e := rows.Scan(&a.ID, &a.Hostname, &a.Enabled, &a.LastIP, &a.FirstSeen, &a.LastSeen); e != nil {
			return nil, e
		}
		out = append(out, a)
	}
	return out, rows.Err()
}

func (s *Store) UpsertDetection(ctx context.Context, d Detection, tenant string) error {
	_, e := s.Pool.Exec(ctx, `INSERT INTO detections(tenant_id,fingerprint,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,first_seen,last_seen)
VALUES($1,$2,$3,$4,'open',$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$11,$12)
ON CONFLICT(tenant_id,fingerprint) DO UPDATE SET
  first_seen=LEAST(detections.first_seen,EXCLUDED.first_seen),
  last_seen=GREATEST(detections.last_seen,EXCLUDED.last_seen),
  window_start=LEAST(detections.window_start,EXCLUDED.window_start),
  window_end=GREATEST(detections.window_end,EXCLUDED.window_end),
  hit_count=GREATEST(detections.hit_count,EXCLUDED.hit_count),
  score=GREATEST(detections.score,EXCLUDED.score),
  summary=EXCLUDED.summary, workstation=EXCLUDED.workstation, updated_at=now()`, tenant, d.Fingerprint, d.RuleName, d.Severity, d.Hostname, d.UserName, d.SourceIP, d.Workstation, d.EventCode, d.Score, d.WindowStart, d.WindowEnd, d.Summary, d.Count)
	return e
}

func (s *Store) ListDetections(ctx context.Context, tenant string, limit int) ([]Detection, error) {
	rows, e := s.Pool.Query(ctx, `SELECT id,fingerprint,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,first_seen,last_seen FROM detections WHERE tenant_id=$1 ORDER BY last_seen DESC LIMIT $2`, tenant, limit)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	var out []Detection
	for rows.Next() {
		var d Detection
		if e := rows.Scan(&d.ID, &d.Fingerprint, &d.RuleName, &d.Severity, &d.Status, &d.Hostname, &d.UserName, &d.SourceIP, &d.Workstation, &d.EventCode, &d.Score, &d.WindowStart, &d.WindowEnd, &d.Summary, &d.Count, &d.FirstSeen, &d.LastSeen); e != nil {
			return nil, e
		}
		out = append(out, d)
	}
	return out, rows.Err()
}
func (s *Store) UpdateDetectionStatus(ctx context.Context, tenant string, id int64, status string) error {
	if status != "open" && status != "investigating" && status != "closed" && status != "false_positive" {
		return fmt.Errorf("invalid status")
	}
	_, e := s.Pool.Exec(ctx, `UPDATE detections SET status=$3, updated_at=now() WHERE tenant_id=$1 AND id=$2`, tenant, id, status)
	return e
}
func (s *Store) DetectionCounts(ctx context.Context, tenant string) (map[string]int64, error) {
	rows, e := s.Pool.Query(ctx, `SELECT severity,count(*) FROM detections WHERE tenant_id=$1 AND status IN ('open','investigating') GROUP BY severity`, tenant)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	out := map[string]int64{}
	for rows.Next() {
		var k string
		var n int64
		if e := rows.Scan(&k, &n); e != nil {
			return nil, e
		}
		out[k] = n
	}
	return out, rows.Err()
}

func hashHex(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
func secureEqual(a, b string) bool {
	if len(a) != len(b) {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
}
func ClientIP(remote string) string {
	h, _, e := net.SplitHostPort(remote)
	if e == nil {
		return h
	}
	return remote
}