package datamodels

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"moss-dev/datahub/api"
	"moss-dev/gpt"
	"moss-dev/prompts"
	"moss-dev/schema_utils"

	"github.com/pingcap/log"
	"github.com/pkg/errors"
	"go.uber.org/zap"
)

type ModelAid struct {
	dh *api.DatahubClient
	ai *gpt.GPT
}

func NewModelAid() *ModelAid { return &ModelAid{} }

func (m *ModelAid) SetDatahubClient(client *api.DatahubClient) { m.dh = client }
func (m *ModelAid) SetGPTClient(client *gpt.GPT)               { m.ai = client }

type ModelSuggestion struct {
	UnclearConcepts []struct {
		Column     string `json:"column"`
		Reason     string `json:"reason"`
		Assumption string `json:"assumption"`
	} `json:"unclearConcepts"`

	IndistinguishableColumns []struct {
		Columns    []string `json:"columns"`
		Reason     string   `json:"reason"`
		Assumption string   `json:"assumption"`
	} `json:"indistinguishableColumns"`

	FurtherOptimizations map[string]string `json:"furtherOptimizations"`
}

func (m *ModelSuggestion) String() string {
	var s bytes.Buffer
	s.WriteString("Unclear concepts:\n")
	for _, uc := range m.UnclearConcepts {
		fmt.Fprintf(&s, "column:%s\n    Reason:%s\n    Assumption:%s\n", uc.Column, uc.Reason, uc.Assumption)
	}
	s.WriteString("\nIndistinguishable Columns:\n")
	for _, uc := range m.IndistinguishableColumns {
		fmt.Fprintf(&s, "column:%+v\n    Reason:%s\n    Assumption:%s\n", uc.Columns, uc.Reason, uc.Assumption)
	}
	s.WriteString("\nFurther Optimizations:\n")
	for key, desc := range m.FurtherOptimizations {
		fmt.Fprintf(&s, "%s:%s\n", key, desc)
	}

	return s.String()
}

func (m *ModelAid) EvaluateSchemaDesign(tbl *schema_utils.TableInfo) (*ModelSuggestion, error) {
	if m.ai == nil {
		return nil, errors.New("Please set the gpt client by caling m.SetGPTClient")
	}
	prompt := fmt.Sprintf(prompts.ModelEvaluationTemplate, tbl.CreateTableStatement(true), tbl.Table)
	predict, err := m.ai.Request(context.Background(), prompt)
	if err != nil {
		return nil, errors.Wrap(err, "fail to request AI")
	}
	log.Debug("AI evaluate schema design", zap.String("result", predict))

	ms := &ModelSuggestion{}
	if err = json.Unmarshal([]byte(predict), &ms); err != nil {
		return nil, err
	}

	return ms, nil
}

type ColumnDenitions struct {
	Columns map[string]string `json:"columns"`
}

func (c *ColumnDenitions) String() string {
	var def bytes.Buffer

	def.WriteString("The definition of columns:\n")
	for cn, desc := range c.Columns {
		fmt.Fprintf(&def, "%s:%s\n", cn, desc)
	}

	return def.String()
}

func (m *ModelAid) GenMetaData(urn string, tbl *schema_utils.TableInfo) (*ColumnDenitions, error) {
	if m.ai == nil {
		return nil, errors.New("Please set the gpt client by caling m.SetGPTClient")
	}
	prompt := fmt.Sprintf(prompts.ModelMetaGenTemplate, tbl.CreateTableStatement(false), tbl.Table)
	predict, err := m.ai.Request(context.Background(), prompt)
	if err != nil {
		return nil, errors.Wrap(err, "fail to request AI")
	}

	colDef := &ColumnDenitions{}
	if err = json.Unmarshal([]byte(predict), &colDef); err != nil {
		log.Error("fail to decode result", zap.String("result", predict))
		return nil, err
	}

	if m.dh != nil {
		for col, desc := range colDef.Columns {
			if err := m.dh.UpdateFieldDescription(col, desc, urn); err != nil {
				return nil, errors.Wrapf(err, "update field %s description", col)
			}
		}
	}

	if err := m.dh.AddTag("", "urn:li:tag:metagen", urn); err != nil {
		log.Warn("fail to add tag", zap.Error(err))
	}

	return colDef, nil
}