package main

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"time"

	_ "modernc.org/sqlite"
)

type Audit struct {
	db *sql.DB
}

const auditSchema = `
CREATE TABLE IF NOT EXISTS audit_log (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    ts          INTEGER NOT NULL,
    request_id  TEXT NOT NULL,
    capability  TEXT NOT NULL,
    args_hash   TEXT NOT NULL,
    exit_code   INTEGER NOT NULL,
    prev_hash   TEXT NOT NULL,
    this_hash   TEXT NOT NULL UNIQUE
);
CREATE INDEX IF NOT EXISTS audit_log_ts ON audit_log(ts);
CREATE INDEX IF NOT EXISTS audit_log_request_id ON audit_log(request_id);

CREATE TABLE IF NOT EXISTS audit_shell_eval (
    audit_id    INTEGER PRIMARY KEY,
    cmd         TEXT NOT NULL,
    cwd         TEXT,
    shell       TEXT NOT NULL,
    stdout_b64  TEXT,
    stderr_b64  TEXT,
    FOREIGN KEY (audit_id) REFERENCES audit_log(id)
);
`

func OpenAudit(path string) (*Audit, error) {
	db, err := sql.Open("sqlite", path+"?_pragma=journal_mode(wal)&_pragma=foreign_keys(on)")
	if err != nil {
		return nil, fmt.Errorf("open audit db %s: %w", path, err)
	}
	if _, err := db.Exec(auditSchema); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("apply audit schema: %w", err)
	}
	return &Audit{db: db}, nil
}

func (a *Audit) Close() error { return a.db.Close() }

// Append registra una invocacion. Devuelve this_hash.
// args raw → SHA256 → args_hash. Hash chain SHA256(prev || canonical).
func (a *Audit) Append(requestID, capability string, args any, exitCode int) (string, error) {
	argsBytes, err := json.Marshal(args)
	if err != nil {
		return "", fmt.Errorf("marshal args: %w", err)
	}
	argsHash := sha256hex(argsBytes)

	var prevHash string
	_ = a.db.QueryRow("SELECT this_hash FROM audit_log ORDER BY id DESC LIMIT 1").Scan(&prevHash)

	ts := time.Now().Unix()
	canonical := fmt.Sprintf("%s|%d|%s|%s|%s|%d", prevHash, ts, requestID, capability, argsHash, exitCode)
	thisHash := sha256hex([]byte(canonical))

	_, err = a.db.Exec(
		`INSERT INTO audit_log (ts, request_id, capability, args_hash, exit_code, prev_hash, this_hash)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
		ts, requestID, capability, argsHash, exitCode, prevHash, thisHash,
	)
	if err != nil {
		return "", fmt.Errorf("insert audit: %w", err)
	}
	return thisHash, nil
}

func sha256hex(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

// ShellEvalRecord is the verbose audit payload for shell.eval invocations.
// stdout / stderr larger than 4KB are gzip+base64 compressed in the DB.
type ShellEvalRecord struct {
	Cmd    string
	CWD    string
	Shell  string
	Stdout string
	Stderr string
}

// AppendVerbose registra una invocacion + payload cleartext (cmd/cwd/shell + stdout/stderr)
// para forense. Devuelve this_hash. Linkado via audit_log.id en audit_shell_eval.
func (a *Audit) AppendVerbose(requestID, capability string, args any, exitCode int, rec ShellEvalRecord) (string, error) {
	argsBytes, err := json.Marshal(args)
	if err != nil {
		return "", fmt.Errorf("marshal args: %w", err)
	}
	argsHash := sha256hex(argsBytes)

	var prevHash string
	_ = a.db.QueryRow("SELECT this_hash FROM audit_log ORDER BY id DESC LIMIT 1").Scan(&prevHash)

	ts := time.Now().Unix()
	canonical := fmt.Sprintf("%s|%d|%s|%s|%s|%d", prevHash, ts, requestID, capability, argsHash, exitCode)
	thisHash := sha256hex([]byte(canonical))

	tx, err := a.db.Begin()
	if err != nil {
		return "", fmt.Errorf("begin tx: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	res, err := tx.Exec(
		`INSERT INTO audit_log (ts, request_id, capability, args_hash, exit_code, prev_hash, this_hash)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
		ts, requestID, capability, argsHash, exitCode, prevHash, thisHash,
	)
	if err != nil {
		return "", fmt.Errorf("insert audit_log: %w", err)
	}
	auditID, err := res.LastInsertId()
	if err != nil {
		return "", fmt.Errorf("lastinsertid: %w", err)
	}

	stdoutEnc, err := encodeMaybeGzip(rec.Stdout)
	if err != nil {
		return "", fmt.Errorf("encode stdout: %w", err)
	}
	stderrEnc, err := encodeMaybeGzip(rec.Stderr)
	if err != nil {
		return "", fmt.Errorf("encode stderr: %w", err)
	}

	if _, err := tx.Exec(
		`INSERT INTO audit_shell_eval (audit_id, cmd, cwd, shell, stdout_b64, stderr_b64)
		 VALUES (?, ?, ?, ?, ?, ?)`,
		auditID, rec.Cmd, rec.CWD, rec.Shell, stdoutEnc, stderrEnc,
	); err != nil {
		return "", fmt.Errorf("insert audit_shell_eval: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return "", fmt.Errorf("commit tx: %w", err)
	}
	return thisHash, nil
}

// encodeMaybeGzip: si len(s) <= 4KB devuelve base64(plain); si supera, gzip+base64.
// Prefijo "gz:" para distinguir. Vacio devuelve "".
func encodeMaybeGzip(s string) (string, error) {
	if s == "" {
		return "", nil
	}
	const threshold = 4096
	if len(s) <= threshold {
		return "plain:" + base64.StdEncoding.EncodeToString([]byte(s)), nil
	}
	var buf bytes.Buffer
	zw := gzip.NewWriter(&buf)
	if _, err := zw.Write([]byte(s)); err != nil {
		_ = zw.Close()
		return "", err
	}
	if err := zw.Close(); err != nil {
		return "", err
	}
	return "gz:" + base64.StdEncoding.EncodeToString(buf.Bytes()), nil
}