package schema

import (
	"encoding/json"
	"math"
	"strconv"
)

const (
	KindCodebase  = "codebase"
	KindKnowledge = "knowledge"
)

// KnowledgeGraph is the AgentTeam Insight graph envelope.
type KnowledgeGraph struct {
	Version string      `json:"version"`
	Kind    string      `json:"kind,omitempty"`
	Project ProjectMeta `json:"project"`
	Nodes   []GraphNode `json:"nodes"`
	Edges   []GraphEdge `json:"edges"`
	Layers  []Layer     `json:"layers"`
	Tour    []TourStep  `json:"tour"`
}

type ProjectMeta struct {
	Name          string   `json:"name"`
	Languages     []string `json:"languages"`
	Frameworks    []string `json:"frameworks"`
	Description   string   `json:"description"`
	AnalyzedAt    string   `json:"analyzedAt"`
	GitCommitHash string   `json:"gitCommitHash"`
}

// GraphNode preserves unknown keys because graph node consumers preserve extension keys.
type GraphNode struct {
	ID            string          `json:"id"`
	Type          string          `json:"type"`
	Name          string          `json:"name"`
	FilePath      string          `json:"filePath,omitempty"`
	LineRange     []int           `json:"lineRange,omitempty"`
	Summary       string          `json:"summary"`
	Tags          []string        `json:"tags"`
	Complexity    string          `json:"complexity"`
	LanguageNotes string          `json:"languageNotes,omitempty"`
	DomainMeta    *DomainMeta     `json:"domainMeta,omitempty"`
	KnowledgeMeta *KnowledgeMeta  `json:"knowledgeMeta,omitempty"`
	Extra         map[string]JSON `json:"-"`
	shapeInvalid  string
}

// JSON aliases json.RawMessage for passthrough fields without exposing encoding details everywhere.
type JSON = json.RawMessage

type graphNodeJSON struct {
	ID            string         `json:"id"`
	Type          string         `json:"type"`
	Name          string         `json:"name"`
	FilePath      string         `json:"filePath,omitempty"`
	LineRange     []int          `json:"lineRange,omitempty"`
	Summary       string         `json:"summary"`
	Tags          []string       `json:"tags"`
	Complexity    string         `json:"complexity"`
	LanguageNotes string         `json:"languageNotes,omitempty"`
	DomainMeta    *DomainMeta    `json:"domainMeta,omitempty"`
	KnowledgeMeta *KnowledgeMeta `json:"knowledgeMeta,omitempty"`
}

func (node GraphNode) MarshalJSON() ([]byte, error) {
	fields := copyExtra(node.Extra)
	tags := node.Tags
	if tags == nil {
		tags = []string{}
	}
	known := graphNodeJSON{
		ID:            node.ID,
		Type:          node.Type,
		Name:          node.Name,
		FilePath:      node.FilePath,
		LineRange:     node.LineRange,
		Summary:       node.Summary,
		Tags:          tags,
		Complexity:    node.Complexity,
		LanguageNotes: node.LanguageNotes,
		DomainMeta:    node.DomainMeta,
		KnowledgeMeta: node.KnowledgeMeta,
	}
	if err := overlayKnownFields(fields, known); err != nil {
		return nil, err
	}
	return json.Marshal(fields)
}

func (node *GraphNode) UnmarshalJSON(data []byte) error {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}
	var known graphNodeJSON
	if err := json.Unmarshal(data, &known); err != nil {
		return err
	}
	deleteKnown(raw, "id", "type", "name", "filePath", "lineRange", "summary", "tags", "complexity", "languageNotes", "domainMeta", "knowledgeMeta")
	*node = GraphNode{
		ID:            known.ID,
		Type:          known.Type,
		Name:          known.Name,
		FilePath:      known.FilePath,
		LineRange:     known.LineRange,
		Summary:       known.Summary,
		Tags:          known.Tags,
		Complexity:    known.Complexity,
		LanguageNotes: known.LanguageNotes,
		DomainMeta:    known.DomainMeta,
		KnowledgeMeta: known.KnowledgeMeta,
		Extra:         nonEmptyExtra(raw),
	}
	return nil
}

// DomainMeta preserves unknown keys because domain metadata consumers preserve extension keys.
type DomainMeta struct {
	Entities                []string        `json:"entities,omitempty"`
	BusinessRules           []string        `json:"businessRules,omitempty"`
	CrossDomainInteractions []string        `json:"crossDomainInteractions,omitempty"`
	EntryPoint              string          `json:"entryPoint,omitempty"`
	EntryType               string          `json:"entryType,omitempty"`
	Extra                   map[string]JSON `json:"-"`
	shapeInvalid            string
}

type domainMetaJSON struct {
	Entities                []string `json:"entities,omitempty"`
	BusinessRules           []string `json:"businessRules,omitempty"`
	CrossDomainInteractions []string `json:"crossDomainInteractions,omitempty"`
	EntryPoint              string   `json:"entryPoint,omitempty"`
	EntryType               string   `json:"entryType,omitempty"`
}

func (meta DomainMeta) MarshalJSON() ([]byte, error) {
	fields := copyExtra(meta.Extra)
	known := domainMetaJSON{
		Entities:                meta.Entities,
		BusinessRules:           meta.BusinessRules,
		CrossDomainInteractions: meta.CrossDomainInteractions,
		EntryPoint:              meta.EntryPoint,
		EntryType:               meta.EntryType,
	}
	if err := overlayKnownFields(fields, known); err != nil {
		return nil, err
	}
	return json.Marshal(fields)
}

func (meta *DomainMeta) UnmarshalJSON(data []byte) error {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}
	var known domainMetaJSON
	if err := json.Unmarshal(data, &known); err != nil {
		return err
	}
	deleteKnown(raw, "entities", "businessRules", "crossDomainInteractions", "entryPoint", "entryType")
	*meta = DomainMeta{
		Entities:                known.Entities,
		BusinessRules:           known.BusinessRules,
		CrossDomainInteractions: known.CrossDomainInteractions,
		EntryPoint:              known.EntryPoint,
		EntryType:               known.EntryType,
		Extra:                   nonEmptyExtra(raw),
	}
	return nil
}

// KnowledgeMeta preserves unknown keys because knowledge metadata consumers preserve extension keys.
type KnowledgeMeta struct {
	Wikilinks    []string        `json:"wikilinks,omitempty"`
	Backlinks    []string        `json:"backlinks,omitempty"`
	Category     string          `json:"category,omitempty"`
	Content      string          `json:"content,omitempty"`
	Extra        map[string]JSON `json:"-"`
	shapeInvalid string
}

type knowledgeMetaJSON struct {
	Wikilinks []string `json:"wikilinks,omitempty"`
	Backlinks []string `json:"backlinks,omitempty"`
	Category  string   `json:"category,omitempty"`
	Content   string   `json:"content,omitempty"`
}

func (meta KnowledgeMeta) MarshalJSON() ([]byte, error) {
	fields := copyExtra(meta.Extra)
	known := knowledgeMetaJSON{
		Wikilinks: meta.Wikilinks,
		Backlinks: meta.Backlinks,
		Category:  meta.Category,
		Content:   meta.Content,
	}
	if err := overlayKnownFields(fields, known); err != nil {
		return nil, err
	}
	return json.Marshal(fields)
}

func (meta *KnowledgeMeta) UnmarshalJSON(data []byte) error {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}
	var known knowledgeMetaJSON
	if err := json.Unmarshal(data, &known); err != nil {
		return err
	}
	deleteKnown(raw, "wikilinks", "backlinks", "category", "content")
	*meta = KnowledgeMeta{
		Wikilinks: known.Wikilinks,
		Backlinks: known.Backlinks,
		Category:  known.Category,
		Content:   known.Content,
		Extra:     nonEmptyExtra(raw),
	}
	return nil
}

// GraphEdge intentionally does not preserve unknown keys because graph edge consumers do not preserve extension keys.
type GraphEdge struct {
	Source        string  `json:"source"`
	Target        string  `json:"target"`
	Type          string  `json:"type"`
	Direction     string  `json:"direction"`
	Description   string  `json:"description,omitempty"`
	Weight        float64 `json:"weight"`
	weightInvalid bool
	shapeInvalid  string
}

type graphEdgeJSON struct {
	Source      string          `json:"source"`
	Target      string          `json:"target"`
	Type        string          `json:"type"`
	Direction   string          `json:"direction"`
	Description string          `json:"description,omitempty"`
	Weight      json.RawMessage `json:"weight"`
}

func (edge *GraphEdge) UnmarshalJSON(data []byte) error {
	var known graphEdgeJSON
	if err := json.Unmarshal(data, &known); err != nil {
		return err
	}
	parsedWeight, weightInvalid, err := parseWeight(known.Weight)
	if err != nil {
		return err
	}
	*edge = GraphEdge{
		Source:        known.Source,
		Target:        known.Target,
		Type:          known.Type,
		Direction:     known.Direction,
		Description:   known.Description,
		Weight:        parsedWeight,
		weightInvalid: weightInvalid,
	}
	return nil
}

func (edge GraphEdge) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		Source      string  `json:"source"`
		Target      string  `json:"target"`
		Type        string  `json:"type"`
		Direction   string  `json:"direction"`
		Description string  `json:"description,omitempty"`
		Weight      float64 `json:"weight"`
	}{
		Source:      edge.Source,
		Target:      edge.Target,
		Type:        edge.Type,
		Direction:   edge.Direction,
		Description: edge.Description,
		Weight:      edge.Weight,
	})
}

func parseWeight(raw json.RawMessage) (float64, bool, error) {
	if len(raw) == 0 || string(raw) == "null" {
		return math.NaN(), false, nil
	}
	var number float64
	if err := json.Unmarshal(raw, &number); err == nil {
		return number, false, nil
	}
	var text string
	if err := json.Unmarshal(raw, &text); err == nil {
		parsed, err := strconv.ParseFloat(text, 64)
		if err == nil {
			return parsed, false, nil
		}
		return math.NaN(), false, nil
	}
	return math.NaN(), true, nil
}

type Layer struct {
	ID           string   `json:"id"`
	Name         string   `json:"name"`
	Description  string   `json:"description"`
	NodeIDs      []string `json:"nodeIds"`
	shapeInvalid string
}

type TourStep struct {
	Order          int      `json:"order"`
	Title          string   `json:"title"`
	Description    string   `json:"description"`
	NodeIDs        []string `json:"nodeIds"`
	LanguageLesson string   `json:"languageLesson,omitempty"`
	shapeInvalid   string
}

func copyExtra(extra map[string]JSON) map[string]json.RawMessage {
	fields := make(map[string]json.RawMessage, len(extra)+8)
	for key, value := range extra {
		copied := append([]byte(nil), value...)
		fields[key] = copied
	}
	return fields
}

func nonEmptyExtra(extra map[string]json.RawMessage) map[string]JSON {
	if len(extra) == 0 {
		return nil
	}
	return extra
}

func deleteKnown(fields map[string]json.RawMessage, keys ...string) {
	for _, key := range keys {
		delete(fields, key)
	}
}

func overlayKnownFields(fields map[string]json.RawMessage, value any) error {
	data, err := json.Marshal(value)
	if err != nil {
		return err
	}
	var known map[string]json.RawMessage
	if err := json.Unmarshal(data, &known); err != nil {
		return err
	}
	for key, raw := range known {
		fields[key] = raw
	}
	return nil
}