// Package devicefleet implements the device MQTT heartbeat and explicit OTA workflow.
package devicefleet

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"io"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
)

type Error struct{ Code, Message string }

func (e *Error) Error() string            { return e.Message }
func conflict(code, message string) error { return &Error{code, message} }

type Store interface {
	FreezeFirmware(context.Context, string, string, string, int64) error
	OpenDownload(context.Context, string) (io.ReadSeekCloser, error)
	PresignUpload(context.Context, string, time.Duration) (string, error)
	PresignMediaDownload(context.Context, string, time.Duration, string, string, bool) (string, error)
	VerifyDigest(context.Context, string, string, int64) error
}
type Service struct {
	Pool            *pgxpool.Pool
	Store           Store
	BrokerURL       string
	DownloadBaseURL string
	Publish         func(string, []byte) error
}
type Release struct {
	ID          string    `json:"id"`
	Model       string    `json:"model"`
	Version     string    `json:"version"`
	Code        int64     `json:"version_code"`
	Notes       string    `json:"release_notes"`
	Size        int64     `json:"size"`
	SHA256      string    `json:"sha256"`
	Status      string    `json:"status"`
	Created     time.Time `json:"created_at"`
	UploadURL   string    `json:"upload_url,omitempty"`
	DownloadURL string    `json:"download_url,omitempty"`
	Object      string    `json:"-"`
}
type ReleaseInput struct {
	Model   string `json:"model"`
	Version string `json:"version"`
	Notes   string `json:"release_notes"`
	Size    int64  `json:"size"`
	SHA256  string `json:"sha256"`
}

const releaseFields = `id::text,model,version,version_code,release_notes,size,sha256,status,created_at,object_ref`

func scanRelease(row interface{ Scan(...any) error }) (r Release, err error) {
	err = row.Scan(&r.ID, &r.Model, &r.Version, &r.Code, &r.Notes, &r.Size, &r.SHA256, &r.Status, &r.Created, &r.Object)
	return
}
func (s *Service) Releases(ctx context.Context, tenant string) ([]Release, error) {
	rows, err := s.Pool.Query(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE tenant_id=$1 ORDER BY version_code DESC LIMIT 200`, tenant)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	result := []Release{}
	for rows.Next() {
		r, e := scanRelease(rows)
		if e != nil {
			return nil, e
		}
		result = append(result, r)
	}
	return result, rows.Err()
}
func (s *Service) Models(ctx context.Context, tenant string) ([]string, error) {
	rows, err := s.Pool.Query(ctx, `SELECT DISTINCT model FROM device WHERE tenant_id=$1 ORDER BY model`, tenant)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	result := []string{}
	for rows.Next() {
		var m string
		if err = rows.Scan(&m); err != nil {
			return nil, err
		}
		result = append(result, m)
	}
	return result, rows.Err()
}
func (s *Service) Prepare(ctx context.Context, tenant string, in ReleaseInput) (Release, error) {
	in.Model = strings.TrimSpace(in.Model)
	in.Version = strings.TrimSpace(in.Version)
	in.SHA256 = strings.ToLower(in.SHA256)
	hash, err := hex.DecodeString(in.SHA256)
	if err != nil || len(hash) != 32 || in.Model == "" || len(in.Model) > 120 || in.Version == "" || len(in.Version) > 80 || len(in.Notes) > 8000 || in.Size <= 0 || in.Size > 32<<20 {
		return Release{}, conflict("INVALID_FIRMWARE", "请选择型号，填写版本名称并上传不超过32MB的固件")
	}
	if s.Store == nil {
		return Release{}, conflict("STORAGE_UNAVAILABLE", "固件存储未配置")
	}
	var exists bool
	if err = s.Pool.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM device WHERE tenant_id=$1 AND model=$2)`, tenant, in.Model).Scan(&exists); err != nil {
		return Release{}, err
	}
	if !exists {
		return Release{}, conflict("MODEL_NOT_FOUND", "请先在设备库存登记该型号")
	}
	// Staging objects are never distributed. Finalize copies a verified immutable object.
	object := "fleet/" + tenant + "/" + uuid.NewString() + ".bin"
	url, err := s.Store.PresignUpload(ctx, object, 15*time.Minute)
	if err != nil {
		return Release{}, err
	}
	r, err := scanRelease(s.Pool.QueryRow(ctx, `INSERT INTO fleet_firmware(tenant_id,model,version,release_notes,size,sha256,object_ref) VALUES($1,$2,$3,$4,$5,$6,$7) ON CONFLICT(tenant_id,model,version) DO UPDATE SET release_notes=excluded.release_notes,size=excluded.size,sha256=excluded.sha256,object_ref=excluded.object_ref WHERE fleet_firmware.status='UPLOADING' RETURNING `+releaseFields, tenant, in.Model, in.Version, in.Notes, in.Size, in.SHA256, object))
	if errors.Is(err, pgx.ErrNoRows) {
		return Release{}, conflict("VERSION_EXISTS", "该型号已有同名已发布版本，请使用新的版本名称")
	}
	r.UploadURL = url
	return r, err
}
func (s *Service) Finalize(ctx context.Context, tenant, id string) (Release, error) {
	r, err := scanRelease(s.Pool.QueryRow(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE tenant_id=$1 AND id=$2`, tenant, id))
	if err != nil {
		return r, err
	}
	if r.Status == "READY" {
		return r, nil
	}
	if s.Store == nil {
		return r, conflict("STORAGE_UNAVAILABLE", "固件存储未配置")
	}
	// SHA-256 is included in every OTA command. Devices must verify it before booting.
	finalObject := "fleet/ready/" + uuid.NewString() + ".bin"
	if err = s.Store.FreezeFirmware(ctx, r.Object, finalObject, r.SHA256, r.Size); err != nil {
		return r, conflict("FIRMWARE_INVALID", "文件大小或SHA256不匹配，请重新上传")
	}
	tag, err := s.Pool.Exec(ctx, `UPDATE fleet_firmware SET status='READY',object_ref=$3 WHERE id=$1 AND tenant_id=$2 AND status='UPLOADING' AND object_ref=$4`, id, tenant, finalObject, r.Object)
	if err != nil {
		return r, err
	}
	if tag.RowsAffected() != 1 {
		return Release{}, conflict("UPLOAD_CHANGED", "上传状态已变化，请刷新后重试")
	}
	return scanRelease(s.Pool.QueryRow(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE id=$1 AND tenant_id=$2`, id, tenant))
}
func (s *Service) Download(ctx context.Context, tenant, id string) (Release, error) {
	r, err := scanRelease(s.Pool.QueryRow(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE tenant_id=$1 AND id=$2 AND status='READY'`, tenant, id))
	if err != nil {
		return r, err
	}
	err = s.download(ctx, &r)
	return r, err
}
func (s *Service) download(ctx context.Context, r *Release) error {
	if s.Store == nil {
		return conflict("STORAGE_UNAVAILABLE", "固件存储未配置")
	}
	_, tokenHash, address, expiresAt, err := newDownloadTicket(s.DownloadBaseURL)
	if err != nil {
		return conflict("DOWNLOAD_UNAVAILABLE", "固件短下载地址未正确配置")
	}
	if _, err = s.Pool.Exec(ctx, `INSERT INTO fleet_firmware_download_ticket(token_hash,release_id,expires_at) VALUES($1,$2,$3)`, tokenHash, r.ID, expiresAt); err != nil {
		return err
	}
	r.DownloadURL = address
	return nil
}
func Password(key []byte, factory string) string {
	mac := hmac.New(sha256.New, key)
	mac.Write([]byte("MQTT1\n" + factory))
	return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}
func StatusTopic(factory string) string { return "vl/devices/" + factory + "/status" }
func OTATopic(factory string) string    { return "vl/devices/" + factory + "/ota" }
func (s *Service) Config(ctx context.Context, device, factory string, key []byte) (map[string]any, error) {
	if s.BrokerURL == "" || s.Publish == nil {
		return nil, conflict("MQTT_UNAVAILABLE", "MQTT服务未启用")
	}
	password := Password(key, factory)
	hash := sha256.Sum256([]byte(password))
	_, err := s.Pool.Exec(ctx, `INSERT INTO fleet_presence(device_id,password_hash) VALUES($1,$2) ON CONFLICT(device_id) DO UPDATE SET password_hash=excluded.password_hash`, device, hash[:])
	if err != nil {
		return nil, err
	}
	return map[string]any{"broker_url": s.BrokerURL, "client_id": factory, "username": factory, "password": password, "heartbeat_interval_seconds": 30, "keepalive_seconds": 60, "status_topic": StatusTopic(factory), "ota_topic": OTATopic(factory), "qos": 1}, nil
}

type Task struct {
	ID      string    `json:"task_id"`
	Status  string    `json:"status"`
	Version string    `json:"version"`
	Code    int64     `json:"version_code"`
	Message string    `json:"message"`
	Created time.Time `json:"created_at"`
}
type State struct {
	Online          bool       `json:"online"`
	LastSeen        *time.Time `json:"last_seen_at"`
	CurrentVersion  string     `json:"current_version"`
	CurrentCode     int64      `json:"current_version_code"`
	CurrentURL      string     `json:"current_download_url"`
	OTASupported    bool       `json:"ota_supported"`
	Latest          *Release   `json:"latest"`
	UpdateAvailable bool       `json:"update_available"`
	Task            *Task      `json:"ota_task"`
}

func (s *Service) State(ctx context.Context, device string) (State, error) {
	var out State
	var tenant, model string
	err := s.Pool.QueryRow(ctx, `SELECT d.tenant_id::text,d.model,COALESCE(p.connected AND p.last_seen_at>now()-interval '90 seconds',false),p.last_seen_at,COALESCE(NULLIF(p.firmware_version,''),dt.firmware_version,''),COALESCE(p.version_code,0),COALESCE(p.ota_supported,false) FROM device d LEFT JOIN fleet_presence p ON p.device_id=d.id LEFT JOIN device_twin dt ON dt.device_id=d.id WHERE d.id=$1`, device).Scan(&tenant, &model, &out.Online, &out.LastSeen, &out.CurrentVersion, &out.CurrentCode, &out.OTASupported)
	if err != nil {
		return out, err
	}
	latest, err := scanRelease(s.Pool.QueryRow(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE tenant_id=$1 AND model=$2 AND status='READY' ORDER BY version_code DESC LIMIT 1`, tenant, model))
	if err == nil {
		if err = s.download(ctx, &latest); err != nil {
			return out, err
		}
		out.Latest = &latest
		out.UpdateAvailable = latest.Code > out.CurrentCode
	} else if !errors.Is(err, pgx.ErrNoRows) {
		return out, err
	}
	current, err := scanRelease(s.Pool.QueryRow(ctx, `SELECT `+releaseFields+` FROM fleet_firmware WHERE tenant_id=$1 AND model=$2 AND version_code=$3 AND status='READY'`, tenant, model, out.CurrentCode))
	if err == nil {
		if err = s.download(ctx, &current); err != nil {
			return out, err
		}
		out.CurrentURL = current.DownloadURL
	} else if !errors.Is(err, pgx.ErrNoRows) {
		return out, err
	}
	var t Task
	err = s.Pool.QueryRow(ctx, `SELECT o.id::text,o.status,f.version,f.version_code,o.message,o.created_at FROM fleet_ota o JOIN fleet_firmware f ON f.id=o.release_id WHERE o.device_id=$1 ORDER BY o.created_at DESC LIMIT 1`, device).Scan(&t.ID, &t.Status, &t.Version, &t.Code, &t.Message, &t.Created)
	if err == nil {
		out.Task = &t
	} else if !errors.Is(err, pgx.ErrNoRows) {
		return out, err
	}
	return out, nil
}

type Report struct {
	TaskID          string `json:"task_id"`
	Status          string `json:"status"`
	FirmwareVersion string `json:"firmware_version"`
	Code            int64  `json:"version_code"`
	Message         string `json:"message,omitempty"`
}

func (s *Service) Report(ctx context.Context, device string, in Report) error {
	if _, err := uuid.Parse(in.TaskID); err != nil {
		return conflict("INVALID_REPORT", "升级任务ID不正确")
	}
	if in.Status != "accepted" && in.Status != "success" && in.Status != "failed" || len(in.Message) > 500 {
		return conflict("INVALID_REPORT", "升级状态不正确")
	}
	tx, err := s.Pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer tx.Rollback(ctx)
	var status, version string
	var code int64
	var created time.Time
	err = tx.QueryRow(ctx, `SELECT o.status,f.version,f.version_code,o.created_at FROM fleet_ota o JOIN fleet_firmware f ON f.id=o.release_id WHERE o.id=$1 AND o.device_id=$2 FOR UPDATE OF o`, in.TaskID, device).Scan(&status, &version, &code, &created)
	if errors.Is(err, pgx.ErrNoRows) {
		return conflict("OTA_NOT_FOUND", "本设备没有该升级任务")
	}
	if err != nil {
		return err
	}
	if in.Status == "success" && (in.Code != code || in.FirmwareVersion != version) {
		return conflict("VERSION_MISMATCH", "升级完成上报必须匹配目标版本和顺序号")
	}
	if status == in.Status {
		return tx.Commit(ctx)
	}
	if in.Status == "accepted" && time.Since(created) > 30*time.Second {
		return conflict("OTA_EXPIRED", "升级指令已过期，请刷新后重新发起")
	}
	if status == "success" || status == "failed" {
		return conflict("OTA_FINISHED", "该升级任务已经结束")
	}
	if status == "timeout" && in.Status == "accepted" {
		return conflict("OTA_EXPIRED", "升级指令已过期")
	}
	_, err = tx.Exec(ctx, `UPDATE fleet_ota SET status=$2,message=$3,updated_at=now() WHERE id=$1`, in.TaskID, in.Status, in.Message)
	if err != nil {
		return err
	}
	if in.Status == "success" || in.Status == "failed" {
		if _, err = tx.Exec(ctx, `UPDATE fleet_firmware_download_ticket SET revoked_at=now() WHERE ota_id=$1 AND revoked_at IS NULL`, in.TaskID); err != nil {
			return err
		}
	}
	if in.Status == "success" {
		var tag pgconn.CommandTag
		tag, err = tx.Exec(ctx, `UPDATE fleet_presence SET firmware_version=$2,version_code=$3 WHERE device_id=$1 AND version_code<=$3`, device, version, code)
		if err != nil {
			return err
		}
		if tag.RowsAffected() == 0 {
			return tx.Commit(ctx)
		}
		_, err = tx.Exec(ctx, `UPDATE device_twin SET firmware_version=$2,updated_at=now() WHERE device_id=$1`, device, version)
		if err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}