package postgres

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

	"example.com/siem-greenfield/internal/rules"
	"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"`
	RuleID      string    `json:"rule_id"`
	RuleSetID   string    `json:"rule_set_id"`
	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"`
	Tags        []string  `json:"tags"`
	MITRE       []string  `json:"mitre"`
}

type RuleSetRecord struct {
	ID           string    `json:"id"`
	Name         string    `json:"name"`
	Description  string    `json:"description"`
	Version      int       `json:"version"`
	Enabled      bool      `json:"enabled"`
	Source       string    `json:"source"`
	Locked       bool      `json:"locked"`
	RuleCount    int64     `json:"rule_count"`
	EnabledRules int64     `json:"enabled_rules"`
	UpdatedAt    time.Time `json:"updated_at"`
}

type Suppression struct {
	ID              int64      `json:"id"`
	RuleID          string     `json:"rule_id"`
	HostPattern     string     `json:"host_pattern"`
	UserPattern     string     `json:"user_pattern"`
	SourceIPPattern string     `json:"source_ip_pattern"`
	Reason          string     `json:"reason"`
	Enabled         bool       `json:"enabled"`
	ExpiresAt       *time.Time `json:"expires_at,omitempty"`
	CreatedAt       time.Time  `json:"created_at"`
}

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
		}
		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) SetAgentEnabled(ctx context.Context, tenant, id string, enabled bool) error {
	tag, e := s.Pool.Exec(ctx, `UPDATE agents SET enabled=$3 WHERE tenant_id=$1 AND id=$2::uuid`, tenant, id, enabled)
	if e != nil {
		return e
	}
	if tag.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}

func (s *Store) SyncRuleSet(ctx context.Context, tenant string, rs rules.RuleSet) error {
	tx, e := s.Pool.Begin(ctx)
	if e != nil {
		return e
	}
	defer tx.Rollback(ctx)
	_, e = tx.Exec(ctx, `INSERT INTO rule_sets(tenant_id,id,name,description,version,enabled,source,locked)
VALUES($1,$2,$3,$4,$5,$6,'builtin',true)
ON CONFLICT(tenant_id,id) DO UPDATE SET name=EXCLUDED.name,description=EXCLUDED.description,version=GREATEST(rule_sets.version,EXCLUDED.version),source='builtin',locked=true,updated_at=now()`, tenant, rs.ID, rs.Name, rs.Description, rs.Version, rs.Enabled)
	if e != nil {
		return e
	}
	ids := make([]string, 0, len(rs.Rules))
	for _, r := range rs.Rules {
		ids = append(ids, r.ID)
		b, e := json.Marshal(r)
		if e != nil {
			return e
		}
		_, e = tx.Exec(ctx, `INSERT INTO detection_rules(tenant_id,id,rule_set_id,title,severity,score,enabled,source,definition)
VALUES($1,$2,$3,$4,$5,$6,$7,'builtin',$8::jsonb)
ON CONFLICT(tenant_id,id) DO UPDATE SET rule_set_id=EXCLUDED.rule_set_id,title=EXCLUDED.title,severity=EXCLUDED.severity,score=EXCLUDED.score,source='builtin',definition=EXCLUDED.definition,updated_at=now()`, tenant, r.ID, rs.ID, r.Title, r.Severity, r.Score, r.Enabled, string(b))
		if e != nil {
			return e
		}
	}
	if len(ids) == 0 {
		_, e = tx.Exec(ctx, `DELETE FROM detection_rules WHERE tenant_id=$1 AND rule_set_id=$2 AND source='builtin'`, tenant, rs.ID)
	} else {
		_, e = tx.Exec(ctx, `DELETE FROM detection_rules WHERE tenant_id=$1 AND rule_set_id=$2 AND source='builtin' AND NOT (id = ANY($3))`, tenant, rs.ID, ids)
	}
	if e != nil {
		return e
	}
	return tx.Commit(ctx)
}

func (s *Store) PruneBuiltinRuleSets(ctx context.Context, tenant string, keep []string) error {
	if len(keep) == 0 {
		_, e := s.Pool.Exec(ctx, `DELETE FROM rule_sets WHERE tenant_id=$1 AND source='builtin'`, tenant)
		return e
	}
	_, e := s.Pool.Exec(ctx, `DELETE FROM rule_sets WHERE tenant_id=$1 AND source='builtin' AND NOT (id = ANY($2))`, tenant, keep)
	return e
}

func (s *Store) EnsureCustomRuleSet(ctx context.Context, tenant string) error {
	_, e := s.Pool.Exec(ctx, `INSERT INTO rule_sets(tenant_id,id,name,description,version,enabled,source,locked)
VALUES($1,'custom','Eigene Regeln','Über die SIEM-Oberfläche verwaltete Regeln',1,true,'custom',false)
ON CONFLICT(tenant_id,id) DO NOTHING`, tenant)
	return e
}

func (s *Store) ListRuleSets(ctx context.Context, tenant string) ([]RuleSetRecord, error) {
	rows, e := s.Pool.Query(ctx, `SELECT rs.id,rs.name,rs.description,rs.version,rs.enabled,rs.source,rs.locked,rs.updated_at,
count(r.id),count(r.id) FILTER (WHERE r.enabled)
FROM rule_sets rs LEFT JOIN detection_rules r ON r.tenant_id=rs.tenant_id AND r.rule_set_id=rs.id
WHERE rs.tenant_id=$1 GROUP BY rs.id,rs.name,rs.description,rs.version,rs.enabled,rs.source,rs.locked,rs.updated_at ORDER BY rs.name`, tenant)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	var out []RuleSetRecord
	for rows.Next() {
		var x RuleSetRecord
		if e := rows.Scan(&x.ID, &x.Name, &x.Description, &x.Version, &x.Enabled, &x.Source, &x.Locked, &x.UpdatedAt, &x.RuleCount, &x.EnabledRules); e != nil {
			return nil, e
		}
		out = append(out, x)
	}
	return out, rows.Err()
}

func (s *Store) ListRules(ctx context.Context, tenant string, enabledOnly bool) ([]rules.StoredRule, error) {
	q := `SELECT r.rule_set_id,rs.name,rs.enabled,r.source,r.enabled,r.definition FROM detection_rules r JOIN rule_sets rs ON rs.tenant_id=r.tenant_id AND rs.id=r.rule_set_id WHERE r.tenant_id=$1`
	if enabledOnly {
		q += ` AND r.enabled AND rs.enabled`
	}
	q += ` ORDER BY rs.name,r.severity DESC,r.title`
	rows, e := s.Pool.Query(ctx, q, tenant)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	var out []rules.StoredRule
	for rows.Next() {
		var x rules.StoredRule
		var b []byte
		var enabled bool
		if e := rows.Scan(&x.RuleSetID, &x.RuleSetName, &x.RuleSetOn, &x.Source, &enabled, &b); e != nil {
			return nil, e
		}
		if e := json.Unmarshal(b, &x.Rule); e != nil {
			return nil, e
		}
		x.Enabled = enabled
		out = append(out, x)
	}
	return out, rows.Err()
}
func (s *Store) SetRuleEnabled(ctx context.Context, tenant, id string, enabled bool) error {
	tag, e := s.Pool.Exec(ctx, `UPDATE detection_rules SET enabled=$3,updated_at=now() WHERE tenant_id=$1 AND id=$2`, tenant, id, enabled)
	if e != nil {
		return e
	}
	if tag.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}
func (s *Store) SetRuleSetEnabled(ctx context.Context, tenant, id string, enabled bool) error {
	tag, e := s.Pool.Exec(ctx, `UPDATE rule_sets SET enabled=$3,updated_at=now() WHERE tenant_id=$1 AND id=$2`, tenant, id, enabled)
	if e != nil {
		return e
	}
	if tag.RowsAffected() == 0 {
		return pgx.ErrNoRows
	}
	return nil
}
func (s *Store) SaveCustomRule(ctx context.Context, tenant string, r rules.Rule) error {
	if err := rules.Validate(r); err != nil {
		return err
	}
	if err := s.EnsureCustomRuleSet(ctx, tenant); err != nil {
		return err
	}
	b, e := json.Marshal(r)
	if e != nil {
		return e
	}
	tag, e := s.Pool.Exec(ctx, `INSERT INTO detection_rules(tenant_id,id,rule_set_id,title,severity,score,enabled,source,definition)
VALUES($1,$2,'custom',$3,$4,$5,$6,'custom',$7::jsonb)
ON CONFLICT(tenant_id,id) DO UPDATE SET rule_set_id='custom',title=EXCLUDED.title,severity=EXCLUDED.severity,score=EXCLUDED.score,enabled=EXCLUDED.enabled,definition=EXCLUDED.definition,updated_at=now()
WHERE detection_rules.source='custom'`, tenant, r.ID, r.Title, r.Severity, r.Score, r.Enabled, string(b))
	if e != nil {
		return e
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("rule id %q belongs to a built-in rule set", r.ID)
	}
	return nil
}

func (s *Store) CreateSuppression(ctx context.Context, tenant string, x Suppression) error {
	if x.RuleID == "" && x.HostPattern == "" && x.UserPattern == "" && x.SourceIPPattern == "" {
		return fmt.Errorf("suppression must constrain rule or entity")
	}
	_, e := s.Pool.Exec(ctx, `INSERT INTO detection_suppressions(tenant_id,rule_id,host_pattern,user_pattern,source_ip_pattern,reason,enabled,expires_at) VALUES($1,$2,$3,$4,$5,$6,true,$7)`, tenant, x.RuleID, x.HostPattern, x.UserPattern, x.SourceIPPattern, x.Reason, x.ExpiresAt)
	return e
}
func (s *Store) DeleteSuppression(ctx context.Context, tenant string, id int64) error {
	_, e := s.Pool.Exec(ctx, `DELETE FROM detection_suppressions WHERE tenant_id=$1 AND id=$2`, tenant, id)
	return e
}
func (s *Store) ListSuppressions(ctx context.Context, tenant string) ([]Suppression, error) {
	rows, e := s.Pool.Query(ctx, `SELECT id,rule_id,host_pattern,user_pattern,source_ip_pattern,reason,enabled,expires_at,created_at FROM detection_suppressions WHERE tenant_id=$1 ORDER BY created_at DESC`, tenant)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	var out []Suppression
	for rows.Next() {
		var x Suppression
		if e := rows.Scan(&x.ID, &x.RuleID, &x.HostPattern, &x.UserPattern, &x.SourceIPPattern, &x.Reason, &x.Enabled, &x.ExpiresAt, &x.CreatedAt); e != nil {
			return nil, e
		}
		out = append(out, x)
	}
	return out, rows.Err()
}
func (s *Store) IsSuppressed(ctx context.Context, tenant, ruleID, host, user, ip string, at time.Time) (bool, error) {
	var ok bool
	e := s.Pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM detection_suppressions WHERE tenant_id=$1 AND enabled AND (expires_at IS NULL OR expires_at>$6)
AND (rule_id='' OR rule_id=$2)
AND (host_pattern='' OR $3 LIKE replace(host_pattern,'*','%'))
AND (user_pattern='' OR $4 LIKE replace(user_pattern,'*','%'))
AND (source_ip_pattern='' OR $5 LIKE replace(source_ip_pattern,'*','%')))`, tenant, ruleID, host, user, ip, at).Scan(&ok)
	return ok, e
}

func (s *Store) UpsertDetection(ctx context.Context, d Detection, tenant string) error {
	tags, _ := json.Marshal(d.Tags)
	mitre, _ := json.Marshal(d.MITRE)
	_, e := s.Pool.Exec(ctx, `INSERT INTO detections(tenant_id,fingerprint,rule_id,rule_set_id,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,tags,mitre,first_seen,last_seen)
VALUES($1,$2,$3,$4,$5,$6,'open',$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17::jsonb,$18::jsonb,$13,$14)
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,tags=EXCLUDED.tags,mitre=EXCLUDED.mitre,updated_at=now()`, tenant, d.Fingerprint, d.RuleID, d.RuleSetID, d.RuleName, d.Severity, d.Hostname, d.UserName, d.SourceIP, d.Workstation, d.EventCode, d.Score, d.WindowStart, d.WindowEnd, d.Summary, d.Count, string(tags), string(mitre))
	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_id,rule_set_id,rule_name,severity,status,hostname,user_name,source_ip,workstation,event_code,score,window_start,window_end,summary,hit_count,tags,mitre,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
		var tags, mitre []byte
		if e := rows.Scan(&d.ID, &d.Fingerprint, &d.RuleID, &d.RuleSetID, &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, &tags, &mitre, &d.FirstSeen, &d.LastSeen); e != nil {
			return nil, e
		}
		_ = json.Unmarshal(tags, &d.Tags)
		_ = json.Unmarshal(mitre, &d.MITRE)
		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
}