package tron

import (
  "crypto/sha256"
  "encoding/hex"
)

// This file implements the minimal protobuf encoding required to compute TRON
// transaction ids (sha256 of the serialized Transaction.raw message) and to
// mirror tronpy's JSON transaction representation. Only the fields the wallet
// populates are handled; proto3 default (zero) values are omitted.

// --- protobuf wire primitives ---

const (
  wireVarint = 0
  wireBytes  = 2
)

type pbWriter struct {
  buf []byte
}

func (w *pbWriter) rawVarint(v uint64) {
  for v >= 0x80 {
    w.buf = append(w.buf, byte(v)|0x80)
    v >>= 7
  }
  w.buf = append(w.buf, byte(v))
}

func (w *pbWriter) tag(field, wire int) {
  w.rawVarint(uint64(field)<<3 | uint64(wire))
}

func (w *pbWriter) varintField(field int, v uint64) {
  if v == 0 {
    return
  }
  w.tag(field, wireVarint)
  w.rawVarint(v)
}

func (w *pbWriter) int64Field(field int, v int64) {
  w.varintField(field, uint64(v))
}

func (w *pbWriter) bytesField(field int, b []byte) {
  if len(b) == 0 {
    return
  }
  w.tag(field, wireBytes)
  w.rawVarint(uint64(len(b)))
  w.buf = append(w.buf, b...)
}

func (w *pbWriter) stringField(field int, s string) {
  w.bytesField(field, []byte(s))
}

func (w *pbWriter) messageField(field int, fn func(*pbWriter)) {
  sub := &pbWriter{}
  fn(sub)
  w.bytesField(field, sub.buf)
}

// --- contract value types (JSON shape matches tronpy dicts) ---

// TransferValue is the value payload of a TransferContract.
type TransferValue struct {
  OwnerAddress string `json:"owner_address"`
  ToAddress    string `json:"to_address"`
  Amount       int64  `json:"amount"`
}

// TriggerValue is the value payload of a TriggerSmartContract.
type TriggerValue struct {
  OwnerAddress    string `json:"owner_address"`
  ContractAddress string `json:"contract_address"`
  Data            string `json:"data"`
  CallTokenValue  int64  `json:"call_token_value"`
  CallValue       int64  `json:"call_value"`
  TokenID         int64  `json:"token_id"`
}

// FreezeValue is the value payload of a FreezeBalanceV2Contract.
type FreezeValue struct {
  OwnerAddress  string `json:"owner_address"`
  FrozenBalance int64  `json:"frozen_balance"`
  Resource      string `json:"resource"`
}

// UnfreezeValue is the value payload of an UnfreezeBalanceV2Contract.
type UnfreezeValue struct {
  OwnerAddress    string `json:"owner_address"`
  UnfreezeBalance int64  `json:"unfreeze_balance"`
  Resource        string `json:"resource"`
}

// ParameterJSON mirrors tronpy's {value, type_url} contract parameter.
type ParameterJSON struct {
  Value   any    `json:"value"`
  TypeURL string `json:"type_url"`
}

// ContractJSON mirrors tronpy's contract entry inside raw_data.
type ContractJSON struct {
  Parameter ParameterJSON `json:"parameter"`
  Type      string        `json:"type"`
}

// RawData mirrors tronpy's raw_data dict; addresses are hex strings.
type RawData struct {
  Contract      []ContractJSON `json:"contract"`
  Timestamp     int64          `json:"timestamp"`
  Expiration    int64          `json:"expiration"`
  RefBlockBytes string         `json:"ref_block_bytes"`
  RefBlockHash  string         `json:"ref_block_hash"`
  FeeLimit      int64          `json:"fee_limit,omitempty"`
}

// --- encoding ---

func contractTypeCode(t string) uint64 {
  switch t {
  case "TransferContract":
    return 1
  case "TriggerSmartContract":
    return 31
  case "FreezeBalanceV2Contract":
    return 54
  case "UnfreezeBalanceV2Contract":
    return 55
  }
  return 0
}

func resourceCode(r string) uint64 {
  if r == "ENERGY" {
    return 1
  }
  return 0
}

func (c *ContractJSON) serializeValue() ([]byte, error) {
  w := &pbWriter{}
  switch v := c.Parameter.Value.(type) {
  case TransferValue:
    owner, err := ToRawAddress(v.OwnerAddress)
    if err != nil {
      return nil, err
    }
    to, err := ToRawAddress(v.ToAddress)
    if err != nil {
      return nil, err
    }
    w.bytesField(1, owner)
    w.bytesField(2, to)
    w.int64Field(3, v.Amount)
  case TriggerValue:
    owner, err := ToRawAddress(v.OwnerAddress)
    if err != nil {
      return nil, err
    }
    contractAddr, err := ToRawAddress(v.ContractAddress)
    if err != nil {
      return nil, err
    }
    data, err := hex.DecodeString(v.Data)
    if err != nil {
      return nil, err
    }
    w.bytesField(1, owner)
    w.bytesField(2, contractAddr)
    w.int64Field(3, v.CallValue)
    w.bytesField(4, data)
    w.int64Field(5, v.CallTokenValue)
    w.int64Field(6, v.TokenID)
  case FreezeValue:
    owner, err := ToRawAddress(v.OwnerAddress)
    if err != nil {
      return nil, err
    }
    w.bytesField(1, owner)
    w.int64Field(2, v.FrozenBalance)
    w.varintField(3, resourceCode(v.Resource))
  case UnfreezeValue:
    owner, err := ToRawAddress(v.OwnerAddress)
    if err != nil {
      return nil, err
    }
    w.bytesField(1, owner)
    w.int64Field(2, v.UnfreezeBalance)
    w.varintField(3, resourceCode(v.Resource))
  default:
    return nil, ErrValidation
  }
  return w.buf, nil
}

// serialize encodes the raw data into the protobuf Transaction.raw message.
func (rd *RawData) serialize() ([]byte, error) {
  w := &pbWriter{}
  rbb, err := hex.DecodeString(rd.RefBlockBytes)
  if err != nil {
    return nil, err
  }
  w.bytesField(1, rbb)
  rbh, err := hex.DecodeString(rd.RefBlockHash)
  if err != nil {
    return nil, err
  }
  w.bytesField(4, rbh)
  w.int64Field(8, rd.Expiration)
  for i := range rd.Contract {
    c := &rd.Contract[i]
    val, err := c.serializeValue()
    if err != nil {
      return nil, err
    }
    typeURL := c.Parameter.TypeURL
    code := contractTypeCode(c.Type)
    w.messageField(11, func(sub *pbWriter) {
      sub.varintField(1, code)
      sub.messageField(2, func(anyW *pbWriter) {
        anyW.stringField(1, typeURL)
        anyW.bytesField(2, val)
      })
    })
  }
  w.int64Field(14, rd.Timestamp)
  w.int64Field(18, rd.FeeLimit)
  return w.buf, nil
}

// TxID computes the transaction id (hex) from the raw data.
func (rd *RawData) TxID() (string, error) {
  b, err := rd.serialize()
  if err != nil {
    return "", err
  }
  h := sha256.Sum256(b)
  return hex.EncodeToString(h[:]), nil
}