package etcd

import (
    "go.etcd.io/etcd/client/v3"
    "context"
    "time"
)

// etcd 클라이언트 설정
type EtcdConfig struct {
    Endpoints []string
    Username  string
    Password  string
    Timeout   time.Duration
}

func NewEtcdClient(config EtcdConfig) (*clientv3.Client, error) {
    return clientv3.New(clientv3.Config{
        Endpoints:   config.Endpoints,
        Username:    config.Username,
        Password:    config.Password,
        DialTimeout: config.Timeout,
    })
}

// KV 저장소 관리자
type KVManager struct {
    client  *clientv3.Client
    ctx     context.Context
    timeout time.Duration
}

func NewKVManager(client *clientv3.Client, timeout time.Duration) *KVManager {
    return &KVManager{
        client:  client,
        ctx:     context.Background(),
        timeout: timeout,
    }
}

// 키-값 설정
func (kv *KVManager) Put(key, value string) error {
    ctx, cancel := context.WithTimeout(kv.ctx, kv.timeout)
    defer cancel()
    
    _, err := kv.client.Put(ctx, key, value)
    return err
}

// 키-값 조회
func (kv *KVManager) Get(key string) (string, error) {
    ctx, cancel := context.WithTimeout(kv.ctx, kv.timeout)
    defer cancel()
    
    resp, err := kv.client.Get(ctx, key)
    if err != nil {
        return "", err
    }
    
    if len(resp.Kvs) == 0 {
        return "", nil
    }
    
    return string(resp.Kvs[0].Value), nil
}

// 프리픽스로 키-값 조회
func (kv *KVManager) GetWithPrefix(prefix string) (map[string]string, error) {
    ctx, cancel := context.WithTimeout(kv.ctx, kv.timeout)
    defer cancel()
    
    resp, err := kv.client.Get(ctx, prefix, clientv3.WithPrefix())
    if err != nil {
        return nil, err
    }
    
    result := make(map[string]string)
    for _, kv := range resp.Kvs {
        result[string(kv.Key)] = string(kv.Value)
    }
    
    return result, nil
}

// 분산 잠금 관리자
type LockManager struct {
    client *clientv3.Client
    ctx    context.Context
}

func NewLockManager(client *clientv3.Client) *LockManager {
    return &LockManager{
        client: client,
        ctx:    context.Background(),
    }
}

// 잠금 획득
func (lm *LockManager) AcquireLock(lockName string, ttl int) (clientv3.LeaseID, error) {
    // 임대 생성
    lease, err := lm.client.Grant(lm.ctx, int64(ttl))
    if err != nil {
        return 0, err
    }

    // 잠금 시도
    _, err = lm.client.Put(lm.ctx, lockName, "", clientv3.WithLease(lease.ID))
    if err != nil {
        return 0, err
    }

    return lease.ID, nil
}

// 잠금 해제
func (lm *LockManager) ReleaseLock(lockName string, leaseID clientv3.LeaseID) error {
    // 임대 취소
    _, err := lm.client.Revoke(lm.ctx, leaseID)
    return err
}

// 감시 관리자
type WatchManager struct {
    client *clientv3.Client
    ctx    context.Context
}

func (wm *WatchManager) Watch(key string, handler func(event *clientv3.Event)) {
    watch := wm.client.Watch(wm.ctx, key)
    
    go func() {
        for resp := range watch {
            for _, event := range resp.Events {
                handler(event)
            }
        }
    }()
}

// 프리픽스 감시
func (wm *WatchManager) WatchPrefix(prefix string, handler func(event *clientv3.Event)) {
    watch := wm.client.Watch(wm.ctx, prefix, clientv3.WithPrefix())
    
    go func() {
        for resp := range watch {
            for _, event := range resp.Events {
                handler(event)
            }
        }
    }()
}

// 리더 선출 관리자
type ElectionManager struct {
    client    *clientv3.Client
    election  *clientv3.Election
    leaderKey string
}

func NewElectionManager(client *clientv3.Client, electionName string) *ElectionManager {
    return &ElectionManager{
        client:    client,
        election:  clientv3.NewElection(client, electionName),
        leaderKey: electionName + "/leader",
    }
}

// 리더 선출 시도
func (em *ElectionManager) Campaign(ctx context.Context, value string) error {
    return em.election.Campaign(ctx, value)
}

// 리더 확인
func (em *ElectionManager) IsLeader(ctx context.Context) (bool, error) {
    resp, err := em.election.Leader(ctx)
    if err != nil {
        return false, err
    }
    
    return string(resp.Kvs[0].Value) == em.leaderKey, nil
}

// 리더십 포기
func (em *ElectionManager) Resign(ctx context.Context) error {
    return em.election.Resign(ctx)
}

type ServiceRegistry struct {
    kv         *KVManager
    serviceDir string
}

type ServiceInfo struct {
    Name     string   `json:"name"`
    Address  string   `json:"address"`
    Port     int      `json:"port"`
    Metadata map[string]string `json:"metadata"`
}

func (sr *ServiceRegistry) Register(service ServiceInfo, ttl int64) error {
    key := sr.serviceDir + "/" + service.Name
    value, err := json.Marshal(service)
    if err != nil {
        return err
    }

    // TTL과 함께 서비스 등록
    lease, err := sr.kv.client.Grant(context.Background(), ttl)
    if err != nil {
        return err
    }

    _, err = sr.kv.client.Put(context.Background(), key, string(value), clientv3.WithLease(lease.ID))
    if err != nil {
        return err
    }

    // Keep-alive 설정
    _, err = sr.kv.client.KeepAlive(context.Background(), lease.ID)
    return err
}

func (sr *ServiceRegistry) Discover(serviceName string) ([]ServiceInfo, error) {
    prefix := sr.serviceDir + "/" + serviceName
    kvs, err := sr.kv.GetWithPrefix(prefix)
    if err != nil {
        return nil, err
    }

    services := make([]ServiceInfo, 0)
    for _, value := range kvs {
        var service ServiceInfo
        if err := json.Unmarshal([]byte(value), &service); err != nil {
            continue
        }
        services = append(services, service)
    }

    return services, nil
}