package main

import (
    "fmt"
    "cmp"
    "slices"
    "maps"
    "strings"
    "strconv"
    "errors"
    "io"
    "mime"
    "github.com/prometheus/common/model"
    "github.com/prometheus/prometheus/model/labels"
    "github.com/prometheus/prometheus/model/exemplar"
    "github.com/prometheus/prometheus/model/textparse"
)

func Escape(src string) string{
    return strings.NewReplacer("\n", `\n`, `"`, `\"`, `\`, `\\`).Replace(src)
}
func SortLabels(ls labels.Labels) labels.Labels{
    ls = ls.Copy()
    slices.SortFunc(ls, func(a, b labels.Label) int{
        return cmp.Compare(a.Name, b.Name)
    })
    return ls
}

type Sample struct{
    labels labels.Labels
    value float64
    suffix string
}
func (self Sample) cmpString() string{
    return SortLabels(self.labels).String()
}
func (self Sample) WriteString(b *strings.Builder){
    firstLabel := 0
    if self.labels.Len() != 0 && self.labels[0].Name == labels.MetricName{
        firstLabel = 1
        b.WriteString(self.labels[0].Value)
    }
    if len(self.labels[firstLabel:]) > 0{
        b.WriteString("{")
        for i, l := range self.labels[firstLabel:]{
            if i != 0{
                b.WriteString(",")
            }
            b.WriteString(l.Name + "=" + `"` + Escape(l.Value) + `"`)
        }
        b.WriteString("}")
    }
    b.WriteString(fmt.Sprintf(" %g%s", self.value, self.suffix))
}
func (self Sample) String() string{
    b := strings.Builder{}
    self.WriteString(&b)
    return b.String()
}

type MetricFamily struct{
    name string
    ty model.MetricType
    unit string
    help string
}
type MetricSampleLabelType int
const (
    MetricSampleLabelTypeUnspec= MetricSampleLabelType(0)
    MetricSampleLabelTypeFloat = MetricSampleLabelType(1)
    MetricSampleLabelTypeStr   = MetricSampleLabelType(2)
)
type MetricSampleLabel struct{
    Name string
    Type MetricSampleLabelType
}
type MetricSampleSet struct{
    Name string
    Labels []MetricSampleLabel
}
func (self MetricFamily) WriteHeader(b *strings.Builder){
    b.WriteString("# TYPE ")
    b.WriteString(self.name)
    b.WriteString(" ")
    b.WriteString(string(self.ty))
    b.WriteString("\n")
    if self.unit != ""{
        b.WriteString("# UNIT ")
        b.WriteString(self.name)
        b.WriteString(" ")
        b.WriteString(self.unit)
        b.WriteString("\n")
    }
    if self.help != ""{
        b.WriteString("# HELP ")
        b.WriteString(self.name)
        b.WriteString(" ")
        b.WriteString(Escape(self.help))
        b.WriteString("\n")
    }
}
func (self MetricFamily) Header() string{
    b := strings.Builder{}
    self.WriteHeader(&b)
    return b.String()
}
func (self MetricFamily) MetricNamesLabels() []MetricSampleSet{
    switch self.ty{
    case model.MetricTypeCounter:
        return []MetricSampleSet{
            {Name: self.name + "_total"},
            {Name: self.name + "_created"}}
    case model.MetricTypeGauge:
        return []MetricSampleSet{{Name: self.name}}
    case model.MetricTypeHistogram:
        return []MetricSampleSet{
            {Name: self.name + "_bucket",
                Labels: []MetricSampleLabel{
                    {Name: "le", Type: MetricSampleLabelTypeFloat},
            }},
            {Name: self.name + "_sum"},
            {Name: self.name + "_created"},
            {Name: self.name + "_count"}}
    case model.MetricTypeGaugeHistogram:
        return []MetricSampleSet{
            {Name: self.name + "_bucket",
                Labels: []MetricSampleLabel{
                    {Name: "le", Type: MetricSampleLabelTypeFloat},
            }},
            {Name: self.name + "_gsum"},
            {Name: self.name + "_gcount"}}
    case model.MetricTypeSummary:
        return []MetricSampleSet{
            {Name: self.name + "_sum"},
            {Name: self.name + "_count"},
            {Name: self.name + "_created"},
            {Name: self.name, Labels: []MetricSampleLabel{{Name: "quantile"}}}}
    case model.MetricTypeInfo:
        return []MetricSampleSet{{Name: self.name + "_info"}}
    case model.MetricTypeStateset:
        return []MetricSampleSet{{
                Name: self.name,
                Labels: []MetricSampleLabel{
                    {Name: self.name, Type: MetricSampleLabelTypeStr},
            }}}
    default:
        return []MetricSampleSet{{Name: self.name}}
    }
}
func (self MetricFamily) MetricNames() []string{
    nl := self.MetricNamesLabels()
    ret := make([]string, 0, len(nl))
    for _, nle := range nl{
        ret = append(ret, nle.Name)
    }
    return ret
}
func (self MetricFamily) CheckMetricName(name string) bool{
    return slices.Contains(self.MetricNames(), name)
}
func (self *MetricFamily) Merge(other MetricFamily){
    // assert(self.name == other.name)
    if self.unit == ""{
        self.unit = other.unit
    }else if other.unit != "" && self.unit != other.unit{
        self.unit = ""
        logger.Warn(fmt.Sprintf(
            "Conflicting metric unit of %q", self.name))
    }
    if self.ty == "" || self.ty == model.MetricTypeUnknown{
        self.ty = other.ty
    }else if !(self.ty == "" || self.ty == model.MetricTypeUnknown) &&
                                                        self.ty != other.ty{
        self.ty = model.MetricTypeUnknown
        logger.Warn(fmt.Sprintf(
            "Conflicting metric types of %q", self.name))
    }
    if self.help == ""{
        self.help = other.help
    }else if len(self.help) < len(other.help){
        self.help = other.help
    }
}

type MetricFamilies struct{
    metas map[string]*MetricFamily
    metamap map[string]*MetricFamily
    samples map[string]map[string]Sample
    meta_conflicts map[string]bool
    metamap_conflicts map[string]bool
    sample_conflicts map[string]bool
}
func EmptyMetricFamilies() MetricFamilies{
    return MetricFamilies{
        metas: map[string]*MetricFamily{},
        metamap: map[string]*MetricFamily{},
        samples: map[string]map[string]Sample{},
        meta_conflicts: map[string]bool{},
        metamap_conflicts: map[string]bool{},
        sample_conflicts: map[string]bool{},
    }
}
func (self MetricFamilies) Clone() MetricFamilies{
    ret := EmptyMetricFamilies()
    ret.meta_conflicts = maps.Clone(self.meta_conflicts)
    ret.metamap_conflicts = maps.Clone(self.metamap_conflicts)
    ret.sample_conflicts = maps.Clone(self.sample_conflicts)
    for _, v := range self.metas{
        ret.AddMeta(*v)
    }
    for n, m := range self.samples{
        dm := ret.samples[n]
        if dm == nil{
            dm = map[string]Sample{}
            ret.samples[n] = dm
        }
        for l, s := range m{
            dm[l] = s
        }
    }
    return ret
}
func (self MetricFamilies) WriteString(b *strings.Builder){
    samples := maps.Clone(self.samples)
    for _, mf := range self.metas{
        // labels, __name__
        ms := map[string]map[string]*[]Sample{}
        for _, mss := range mf.MetricNamesLabels(){
            for _, s := range samples[mss.Name]{
                ml := labels.Labels{}
                for _, l := range s.labels{
                    if l.Name == "__name__" ||
                        slices.ContainsFunc(mss.Labels,
                            func(ld MetricSampleLabel)bool{
                                return ld.Name == l.Name
                    }){
                        continue
                    }
                    ml = append(ml, l)
                }
                mls := SortLabels(ml).String()
                ss := ms[mls][mss.Name]
                if ss == nil{
                    if ms[mls] == nil{
                        ms[mls] = map[string]*[]Sample{}
                    }
                    ss = &[]Sample{}
                    ms[mls][mss.Name] = ss
                }
                *ss = append(*ss, s)
            }
            delete(samples, mss.Name)
        }
        if len(ms) == 0{
            continue
        }
        mf.WriteHeader(b)
        for _, s := range ms{
            for _, mss := range mf.MetricNamesLabels(){
                if s[mss.Name] == nil{
                    continue
                }
                slices.SortFunc(*s[mss.Name], func(a,b Sample)int{
                    for _, l := range mss.Labels{
                        as := a.labels.Get(l.Name)
                        bs := b.labels.Get(l.Name)
                        switch l.Type{
                        case MetricSampleLabelTypeUnspec:
                            continue
                        case MetricSampleLabelTypeFloat:
                            an, ae := strconv.ParseFloat(as, 64)
                            bn, be := strconv.ParseFloat(bs, 64)
                            if ae == nil && be == nil{
                                if an < bn{ return -1
                                }else if bn < an{ return 1
                                }else{ continue }
                            }else{
                                if ae == nil{ return 1
                                }else if be == nil{ return -1
                                }else{
                                    if as < bs{ return -1
                                    }else if bs < as{ return 1
                                    }else{ continue }
                                }
                            }
                        case MetricSampleLabelTypeStr:
                            if as < bs{ return -1
                            }else if bs < as{ return 1
                            }else{ continue }
                        }
                    }
                    return 0
                })
                for _, m := range *s[mss.Name]{
                    m.WriteString(b)
                    b.WriteString("\n")
                }
            }
        }
    }
    for _, mn := range samples{
        for _, m := range mn{
            m.WriteString(b)
            b.WriteString("\n")
        }
    }
    b.WriteString("# EOF\n")
}
func (self MetricFamilies) String() string{
    b := strings.Builder{}
    self.WriteString(&b)
    return b.String()
}
func (self *MetricFamilies) addSampleWithKeys(s Sample, mn string, sk string){
    if self.samples[mn] == nil{
        self.samples[mn] = map[string]Sample{}
    }
    if !self.sample_conflicts[sk]{
        _, found := self.samples[mn][sk]
        if found{
            logger.Error(fmt.Sprintf("Conflicting labelset %s", sk))
            self.sample_conflicts[sk] = true
            delete(self.samples[mn], sk)
            if len(self.samples[mn]) == 0{
                delete(self.samples, mn)
            }
        }else{
            if len(s.labels) != 0 && s.labels[0].Name != labels.MetricName &&
                                s.labels.Has(labels.MetricName){
                s.labels = append(labels.New(
                        labels.Label{
                            Name: labels.MetricName,
                            Value: mn,
                        },
                    ),
                    s.labels.DropMetricName()...)
            }
            self.samples[mn][sk] = s
        }
    }
}
func (self *MetricFamilies) AddSample(s Sample){
    self.addSampleWithKeys(s, s.labels.Get(labels.MetricName), s.cmpString())
}
func (self *MetricFamilies) AddMeta(m MetricFamily){
    mn := m.MetricNames()
    if self.meta_conflicts[m.name]{
        return
    }
    for _, mni := range mn{
        if self.metamap_conflicts[mni]{
            return
        }
    }

    c := make([]*MetricFamily, 0, 2)
    if ci := self.metas[m.name]; ci != nil{
        if m.ty == ci.ty{
            ci.Merge(m)
            return
        }else{
            c = append(c, ci)
        }
    }
    for _, mni := range mn{
        if ci := self.metamap[mni]; ci != nil{
            c = append(c, ci)
        }
    }
    if len(c) == 0{
        self.metas[m.name] = &m
        for _, mni := range mn{
            self.metamap[mni] = &m
        }
    }else{
        logger.Warn(fmt.Sprintf(
            "Conflicting metric families %q (%s) and %q (%s)",
            m.name, m.ty, c[0].name, c[0].ty))
        c = append(c, &m)
        for _, ci := range c{
            delete(self.metas, ci.name)
            self.meta_conflicts[ci.name] = true
            for _, mni := range ci.MetricNames(){
                delete(self.metamap, mni)
                self.metamap_conflicts[mni] = true
            }
        }
    }
}
func (self *MetricFamilies) Merge(other MetricFamilies){
    for _, f := range other.metas{
        self.AddMeta(*f)
    }
    for n, m := range other.samples{
        for l, s := range m{
            self.addSampleWithKeys(s, n, l)
        }
    }
}
func (self MetricFamilies) Relabel(
    f func(labels.Labels) (labels.Labels, bool),
) MetricFamilies{
    ret := EmptyMetricFamilies()
    for on, m := range self.metas{
        l := labels.New(
            labels.Label{
                Name: labels.MetricName,
                Value: on,
            },
        )
        l, keep := f(l)
        if keep{
            nm := *m
            nm.name = l.Get(labels.MetricName)
            ret.AddMeta(nm)
        }
    }

    for _, m := range self.samples{
        for _, s := range m{
            l, keep := f(s.labels)
            if keep{
                ns := s
                ns.labels = l
                ret.AddSample(ns)
            }
        }
    }

    return ret
}

func ParseMetrics(text []byte, fty string) (MetricFamilies, error){
    ty, _, err := mime.ParseMediaType(fty)
    if err != nil{
        ty = ""
    }

    // ignore errors
    p, _ := textparse.New(text, fty, false, labels.NewSymbolTable())

    ret := EmptyMetricFamilies()
    metas := map[string]*MetricFamily{}
    curf := (*MetricFamily)(nil)
    byname := func(name string) *MetricFamily{
        if curf == nil || curf.name != name{
            curf = metas[name]
            if curf == nil{
                curf = new(MetricFamily)
                metas[name] = curf
                curf.name = name
            }
        }
        return curf
    }

    for{
        rt, err := p.Next()
        if errors.Is(err, io.EOF){
            break
        }else if err != nil{
            return MetricFamilies{}, err
        }

        switch rt{
        case textparse.EntryType:
            name, ty := p.Type()
            meta := byname(string(name))
            if meta.ty != ""{
                logger.Warn(fmt.Sprintf(
                    "Multiple type header for metric %q", name))
            }
            meta.ty = ty
        case textparse.EntryUnit:
            name, unit := p.Unit()
            meta := byname(string(name))
            if meta.unit != ""{
                logger.Warn(fmt.Sprintf(
                    "Multiple type header for metric %q", name))
            }
            meta.unit = string(unit)
        case textparse.EntryHelp:
            name, help := p.Help()
            meta := byname(string(name))
            if meta.unit != ""{
                logger.Warn(fmt.Sprintf(
                    "Multiple help header for metric %q", name))
                meta.unit += "\n"
            }
            meta.help += string(help)
        case textparse.EntrySeries:
            _, time, val := p.Series()
            if time != nil{
                return MetricFamilies{}, fmt.Errorf("Metric with timestamp")
            }
            l := labels.New()
            p.Metric(&l)
            s := Sample{
                labels: l,
                value: val,
            }
            ex := exemplar.Exemplar{}
            if p.Exemplar(&ex){
                s.suffix = fmt.Sprintf(" # %s %g", ex.Labels.String(), ex.Value)
                if ex.HasTs{
                    s.suffix += fmt.Sprintf(" %.3f", float64(ex.Ts) / 1000.0)
                }
                if p.Exemplar(&ex){
                    logger.Warn("Discarding exemplar")
                }
            }
            ret.AddSample(s)
        case textparse.EntryHistogram:
            // TODO p.Histogram()
            logger.Error("Histogram is not supported")
        case textparse.EntryComment:
            // ignore comments
        //case textparse.EntryInvalid:
        default:
            logger.Error("Invalid entry found")
        }
    }

    for _, m := range metas{
        if m.ty == ""{
            m.ty = model.MetricTypeUnknown
        }
        if m.ty == "counter" && ty != "application/openmetrics-text"{
            if m.name == "go_memstats_alloc_bytes_total"{
                continue    // to avoid warning
            }
            m.name = strings.TrimSuffix(m.name, "_total")
        }
        ret.AddMeta(*m)
    }
    return ret, nil
}