package ontology

import (
	"context"
	"encoding/json"
	"fmt"

	"github.com/jackc/pgx/v5/pgxpool"
)

type PostgresStore struct {
	pool *pgxpool.Pool
}

func NewPostgresStore(pool *pgxpool.Pool) *PostgresStore {
	return &PostgresStore{pool: pool}
}

func toJSON(v any) string {
	b, _ := json.Marshal(v)
	return string(b)
}

func (s *PostgresStore) Save(material Material, concepts []ConceptCandidate, edges []EdgeCandidate, gaps []Gap) error {
	ctx := context.Background()
	_, err := s.pool.Exec(ctx,
		`INSERT INTO ontology_materials (id, title, source_type, body, created_at) VALUES ($1, $2, $3, $4, $5)`,
		material.ID, material.Title, material.SourceType, material.Body, material.CreatedAt)
	if err != nil {
		return fmt.Errorf("insert material: %w", err)
	}

	for _, c := range concepts {
		_, err := s.pool.Exec(ctx,
			`INSERT INTO ontology_concepts (id, material_id, concept_id, concept_label, summary, review_state, evidence, created_at)
			 VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8)`,
			c.ID, material.ID, c.Concept.ID, c.Concept.Label, c.Summary, string(c.ReviewState),
			toJSON(c.Evidence), c.CreatedAt)
		if err != nil {
			return fmt.Errorf("insert concept: %w", err)
		}
	}

	for _, e := range edges {
		_, err := s.pool.Exec(ctx,
			`INSERT INTO ontology_edges (id, from_concept_id, to_concept_id, kind, evidence, created_at)
			 VALUES ($1, $2, $3, $4, $5::jsonb, $6)`,
			e.ID, e.From.ID, e.To.ID, string(e.Kind),
			toJSON(e.Evidence), e.CreatedAt)
		if err != nil {
			return fmt.Errorf("insert edge: %w", err)
		}
	}

	for _, g := range gaps {
		_, err := s.pool.Exec(ctx,
			`INSERT INTO ontology_gaps (id, concept_id, reason, evidence, created_at)
			 VALUES ($1, $2, $3, $4::jsonb, $5)`,
			g.ID, g.Concept.ID, g.Reason,
			toJSON(g.SupportingEvidence), g.CreatedAt)
		if err != nil {
			return fmt.Errorf("insert gap: %w", err)
		}
	}
	return nil
}

func (s *PostgresStore) Snapshot() Snapshot {
	ctx := context.Background()
	var snap Snapshot

	matRows, _ := s.pool.Query(ctx, `SELECT id, title, source_type, body, created_at FROM ontology_materials`)
	defer matRows.Close()
	for matRows.Next() {
		var m Material
		matRows.Scan(&m.ID, &m.Title, &m.SourceType, &m.Body, &m.CreatedAt)
		snap.Materials = append(snap.Materials, m)
	}

	cRows, _ := s.pool.Query(ctx, `SELECT id, concept_id, concept_label, summary, review_state, evidence, created_at FROM ontology_concepts`)
	defer cRows.Close()
	for cRows.Next() {
		var c ConceptCandidate
		var evidenceJSON string
		cRows.Scan(&c.ID, &c.Concept.ID, &c.Concept.Label, &c.Summary, &c.ReviewState, &evidenceJSON, &c.CreatedAt)
		json.Unmarshal([]byte(evidenceJSON), &c.Evidence)
		snap.Concepts = append(snap.Concepts, c)
	}

	eRows, _ := s.pool.Query(ctx, `SELECT id, from_concept_id, to_concept_id, kind, evidence, created_at FROM ontology_edges`)
	defer eRows.Close()
	for eRows.Next() {
		var e EdgeCandidate
		var evidenceJSON string
		eRows.Scan(&e.ID, &e.From.ID, &e.To.ID, &e.Kind, &evidenceJSON, &e.CreatedAt)
		json.Unmarshal([]byte(evidenceJSON), &e.Evidence)
		snap.Edges = append(snap.Edges, e)
	}

	gRows, _ := s.pool.Query(ctx, `SELECT id, concept_id, reason, evidence, created_at FROM ontology_gaps`)
	defer gRows.Close()
	for gRows.Next() {
		var g Gap
		var evidenceJSON string
		gRows.Scan(&g.ID, &g.Concept.ID, &g.Reason, &evidenceJSON, &g.CreatedAt)
		json.Unmarshal([]byte(evidenceJSON), &g.SupportingEvidence)
		snap.Gaps = append(snap.Gaps, g)
	}

	return snap
}