// spatiald — the PostGIS query service for the trip planner.
//
// Native spatial queries over the osm2pgsql-loaded OSM extracts:
//
//	GET /health                      { ok, extracts: {name: poi_count} }
//	GET /near?lat=&lng=&r=           ST_DWithin radius search
//	GET /nearest?lat=&lng=           KNN (ORDER BY geom <-> point)
//	GET /corridor?points=&r=         corridor buffering along a route
//
// Optional filters on all three:
//
//	extract=nh          restrict to one loaded extract
//	kind=amenity=restaurant   exact kind match, or one of the bare values
//	                         (restaurant, museum, cafe, …)
//	name=café           ILIKE substring on the POI name
//	limit=20            default 20, max 200
//
// /corridor's `points` is "lng,lat;lng,lat;…" (a router polyline, downsampled
// is fine — the buffer does the smoothing work).
package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	_ "github.com/lib/pq"
)

var (
	dsn  *sql.DB
	addr string
	// /semantic: embed the query with the local zembed model, cosine-search poi_vec
	embedBase  = getenvDefault("EMBED_BASE", "http://192.168.3.7:1234/v1")
	embedModel = getenvDefault("EMBED_MODEL", "zembed-1-Q4_K_M")
)

func getenvDefault(k, def string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return def
}

// embedOne asks the llama.cpp OpenAI-compatible /embeddings endpoint for a
// single text. Zembed uses a "query: " prompt prefix for search queries.
func embedOne(text string) ([]float32, error) {
	body, _ := json.Marshal(map[string]any{"model": embedModel, "input": []string{"query: " + text}})
	req, err := http.NewRequest("POST", embedBase+"/embeddings", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{Timeout: 60 * time.Second}
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		b, _ := io.ReadAll(io.LimitReader(res.Body, 300))
		return nil, fmt.Errorf("embed http %d: %s", res.StatusCode, b)
	}
	var j struct {
		Data []struct {
			Embedding []float32 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.NewDecoder(res.Body).Decode(&j); err != nil {
		return nil, err
	}
	if len(j.Data) == 0 || len(j.Data[0].Embedding) == 0 {
		return nil, fmt.Errorf("no embedding in response")
	}
	return j.Data[0].Embedding, nil
}

// vecLiteral renders [0.1, -0.2] as the pgvector text literal [0.1,-0.2]
func vecLiteral(v []float32) string {
	parts := make([]string, len(v))
	for i, f := range v {
		parts[i] = strconv.FormatFloat(float64(f), 'g', 6, 32)
	}
	return "[" + strings.Join(parts, ",") + "]"
}

func handleSemantic(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	text := strings.TrimSpace(first(q, "q"))
	if text == "" {
		jerr(w, 400, "q is required")
		return
	}
	vec, err := embedOne(text)
	if err != nil {
		jerr(w, 502, "embedding failed: "+err.Error())
		return
	}
	// lib/pq has no []float32 → vector codec: send the literal "[f1,f2,…]"
	// as text and cast with ::vector
	where, args := filters(q) // extract/kind/name filters still apply
	args = append(args, vecLiteral(vec))
	vecTok := "@@" + strconv.Itoa(len(args)) + "@@"
	args = append(args, limitOf(q))
	limTok := "@@" + strconv.Itoa(len(args)) + "@@"
	// poi_vec carries (extract, osm_type, osm_id, name, kind, vec); the JOIN
	// back to poi yields the geometry (for map pins) and the tag facts.
	query := `SELECT p.extract, p.name, p.kind,
           ST_Y(p.geom) AS lat, ST_X(p.geom) AS lng,
           p.opening_hours, p.fee, p.website,
           1 - (v.vec <=> ` + vecTok + `::vector)::float8 AS score
      FROM poi_vec v
      JOIN poi p ON p.extract = v.extract AND p.osm_type = v.osm_type AND p.osm_id = v.osm_id` + where + `
  ORDER BY v.vec <=> ` + vecTok + `::vector
 LIMIT ` + limTok
	query = renumber(query, 0) // no body placeholders here; just resolve tokens
	rows, err := dsn.Query(query, args...)
	if err != nil {
		jerr(w, 502, "query: "+err.Error())
		return
	}
	var out []map[string]any
	for rows.Next() {
		var extract, name, kind string
		var lat, lng float64
		var oh, fee, web *string
		var score float64
		if err := rows.Scan(&extract, &name, &kind, &lat, &lng, &oh, &fee, &web, &score); err != nil {
			rows.Close()
			jerr(w, 502, err.Error())
			return
		}
		m := map[string]any{"extract": extract, "name": name, "kind": kind,
			"lat": lat, "lng": lng, "score": float64(int64(score*10000)) / 10000}
		if oh != nil {
			m["opening_hours"] = *oh
		}
		if fee != nil {
			m["fee"] = *fee
		}
		if web != nil {
			m["website"] = *web
		}
		out = append(out, m)
	}
	rows.Close()
	if out == nil {
		out = []map[string]any{}
	}
	cors(w)
	b, _ := json.Marshal(out)
	fmt.Fprintf(w, `{"count":%d,"results":`, len(out))
	w.Write(b)
	w.Write([]byte("}"))
}

type poi struct {
	Extract  string  `json:"extract"`
	OsmID    int64   `json:"osm_id"`
	OsmType  string  `json:"osm_type"`
	Name     string  `json:"name"`
	Kind     string  `json:"kind"`
	OpenH    *string `json:"opening_hours,omitempty"`
	Fee      *string `json:"fee,omitempty"`
	Website  *string `json:"website,omitempty"`
	AddrCity *string `json:"addr_city,omitempty"`
	Lat      float64 `json:"lat"`
	Lng      float64 `json:"lng"`
	DistM    float64 `json:"dist_m"`
}

func main() {
	dsnFlag := flag.String("dsn",
		"host=localhost port=5432 user=trips password=trips dbname=trips sslmode=disable",
		"PostgreSQL/PostGIS DSN")
	addrFlag := flag.String("addr", ":5005", "listen address")
	flag.Parse()
	if env := os.Getenv("SPATIAL_DSN"); env != "" {
		*dsnFlag = env // docker-compose sets the in-network DSN
	}
	addr = *addrFlag

	var err error
	dsn, err = sql.Open("postgres", *dsnFlag)
	if err != nil {
		log.Fatalf("sql.Open: %v", err)
	}
	dsn.SetMaxOpenConns(8)
	dsn.SetConnMaxLifetime(30 * time.Minute)
	if err = dsn.Ping(); err != nil {
		log.Printf("spatiald: DB not reachable yet (%v) — serving /health as down until it is", err)
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/health", handleHealth)
	mux.HandleFunc("/near", handleNear)
	mux.HandleFunc("/nearest", handleNearest)
	mux.HandleFunc("/corridor", handleCorridor)
	mux.HandleFunc("/semantic", handleSemantic)
	log.Printf("spatiald: listening on %s (dsn: %s)", addr, *dsnFlag)
	srv := &http.Server{Addr: addr, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
	if err = srv.ListenAndServe(); err != nil {
		log.Fatal(err)
	}
}

func cors(w http.ResponseWriter) {
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Content-Type", "application/json")
}

func jerr(w http.ResponseWriter, code int, msg string) {
	cors(w)
	w.WriteHeader(code)
	fmt.Fprintf(w, `{"error":%q}`, strings.ReplaceAll(msg, `"`, `\"`))
}

// common filters → (where clause, args). kind accepts "amenity=restaurant"
// (exact) or "restaurant" (match the tag value across all kind families).
// filters → (WHERE clause with @@i@@ placeholder tokens, args). Tokens are
// resolved to $i by renumber(); the query body's own $1..$4 placeholders
// would otherwise collide with the filter params.
func filters(q map[string][]string) (string, []any) {
	var conds []string
	var args []any
	param := func(arg any) string {
		args = append(args, arg)
		return "@@" + strconv.Itoa(len(args)) + "@@"
	}
	if e := first(q, "extract"); e != "" {
		conds = append(conds, "poi.extract = "+param(e))
	}
	if k := first(q, "kind"); k != "" {
		if strings.Contains(k, "=") {
			conds = append(conds, "poi.kind = "+param(k))
		} else if isFamily(k) {
			// bare family name: the whole family ("tourism", "leisure", …)
			conds = append(conds, "poi.kind LIKE '"+k+"=%'")
		} else {
			// bare tag value: exact match in any kind family
			v := param(k)
			var parts []string
			for _, fam := range []string{"amenity", "tourism", "shop", "leisure", "historic", "place"} {
				parts = append(parts, "poi.kind = '"+fam+"='||"+v)
			}
			conds = append(conds, "("+strings.Join(parts, " OR ")+")")
		}
	}
	if n := first(q, "name"); n != "" {
		conds = append(conds, "poi.name ILIKE "+param("%"+n+"%"))
	}
	if len(conds) == 0 {
		return "", args
	}
	return " WHERE " + strings.Join(conds, " AND "), args
}
func isFamily(k string) bool {
	for _, f := range []string{"amenity", "tourism", "shop", "leisure", "historic", "place"} {
		if k == f {
			return true
		}
	}
	return false
}

func first(q map[string][]string, k string) string {
	if v := q[k]; len(v) > 0 {
		return strings.TrimSpace(v[0])
	}
	return ""
}

func limitOf(q map[string][]string) int {
	l, err := strconv.Atoi(first(q, "limit"))
	if err != nil || l < 1 {
		return 20
	}
	if l > 200 {
		return 200
	}
	return l
}

// rows → pois; the SELECT must end with a distance column in metres.
func scanPois(rs *sql.Rows) ([]poi, error) {
	var out []poi
	for rs.Next() {
		var p poi
		if err := rs.Scan(&p.Extract, &p.OsmID, &p.OsmType, &p.Name, &p.Kind,
			&p.OpenH, &p.Fee, &p.Website, &p.AddrCity, &p.Lat, &p.Lng, &p.DistM); err != nil {
			return nil, err
		}
		out = append(out, p)
	}
	return out, rs.Err()
}

func writePois(w http.ResponseWriter, pois []poi) {
	if pois == nil {
		pois = []poi{}
	}
	cors(w)
	fmt.Fprintf(w, `{"count":%d,"results":`, len(pois))
	b, _ := json.Marshal(pois)
	w.Write(b)
	w.Write([]byte("}"))
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	cors(w)
	type ext struct {
		Name     string  `json:"name"`
		PoiCount *int64  `json:"poi_count"`
		Imported *string `json:"imported_at,omitempty"`
	}
	healthy := dsn.Ping() == nil
	var extracts []ext
	if healthy {
		rows, err := dsn.Query(`SELECT name, poi_count, to_char(imported_at, 'YYYY-MM-DD"T"HH24:MI:SSOF') FROM spatial_extract ORDER BY name`)
		if err == nil {
			for rows.Next() {
				var e ext
				rows.Scan(&e.Name, &e.PoiCount, &e.Imported)
				extracts = append(extracts, e)
			}
			rows.Close()
		}
	}
	if extracts == nil {
		extracts = []ext{}
	}
	fmt.Fprintf(w, `{"ok":%v,"postgis":%v,"extracts":`, healthy, healthy)
	b, _ := json.Marshal(extracts)
	w.Write(b)
	w.Write([]byte("}"))
}

// +lat/lng so the client can drop map pins on the results
const poiCols = `poi.extract, poi.osm_id, poi.osm_type, poi.name, poi.kind,
       poi.opening_hours, poi.fee, poi.website, poi.addr_city,
       ST_Y(poi.geom) AS lat, ST_X(poi.geom) AS lng`

func handleNear(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	lat, lng, err := point(q)
	if err != nil {
		jerr(w, 400, err.Error())
		return
	}
	rad, err := radius(q)
	if err != nil {
		jerr(w, 400, err.Error())
		return
	}
	where, args := filters(q)
	args = append(args, lat, lng, rad, limitOf(q))
	// point geometry first, then geography cast for the radius test, then KNN
	// ordering among hits — ST_DWithin(geography) uses the GIST index.
	query := `SELECT ` + poiCols + `,
           ST_Distance(poi.geom::geography, ST_SetSRID(ST_MakePoint($2, $1), 4326)::geography) AS dist_m
      FROM poi` + where + `
     AND ST_DWithin(poi.geom::geography, ST_SetSRID(ST_MakePoint($2, $1), 4326)::geography, $3)
  ORDER BY poi.geom <-> ST_SetSRID(ST_MakePoint($2, $1), 4326)
 LIMIT $4`
	query = renumber(query, len(args)-4)
	rows, err := dsn.Query(query, args...)
	if err != nil {
		jerr(w, 502, "query: "+err.Error())
		return
	}
	pois, err := scanPois(rows)
	rows.Close()
	if err != nil {
		jerr(w, 502, err.Error())
		return
	}
	writePois(w, pois)
}

func handleNearest(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	lat, lng, err := point(q)
	if err != nil {
		jerr(w, 400, err.Error())
		return
	}
	where, args := filters(q)
	args = append(args, lat, lng, limitOf(q))
	query := `SELECT ` + poiCols + `,
           ST_Distance(poi.geom::geography, ST_SetSRID(ST_MakePoint($2, $1), 4326)::geography) AS dist_m
      FROM poi` + where + `
  ORDER BY poi.geom <-> ST_SetSRID(ST_MakePoint($2, $1), 4326)
 LIMIT $3`
	query = renumber(query, len(args)-3)
	rows, err := dsn.Query(query, args...)
	if err != nil {
		jerr(w, 502, "query: "+err.Error())
		return
	}
	pois, err := scanPois(rows)
	rows.Close()
	if err != nil {
		jerr(w, 502, err.Error())
		return
	}
	writePois(w, pois)
}

func handleCorridor(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	rad, err := radius(q)
	if err != nil {
		jerr(w, 400, err.Error())
		return
	}
	pts := parsePoints(first(q, "points"))
	if len(pts) < 2 {
		jerr(w, 400, "points must be 'lng,lat;lng,lat;…' with at least 2 points")
		return
	}
	where, args := filters(q)
	var coords []string
	for _, p := range pts {
		coords = append(coords, fmt.Sprintf("%s %s", strconv.FormatFloat(p[0], 'f', 6, 64), strconv.FormatFloat(p[1], 'f', 6, 64))) // WKT: lng lat
	}
	line := "LINESTRING(" + strings.Join(coords, ",") + ")"
	args = append(args, line, rad, limitOf(q))
	// distance to a line uses ST_Distance(geography) (ST_Distance_Sphere is
	// point-to-point only); the corridor is a buffered geography band.
	query := `SELECT ` + poiCols + `,
           ST_Distance(poi.geom::geography, ST_SetSRID(ST_GeomFromText($1), 4326)::geography) AS dist_m
      FROM poi` + where + `
     AND ST_DWithin(poi.geom::geography,
           ST_Buffer(ST_SetSRID(ST_GeomFromText($1), 4326)::geography, $2), 0, true)
  ORDER BY dist_m
 LIMIT $3`
	query = renumber(query, len(args)-3)
	rows, err := dsn.Query(query, args...)
	if err != nil {
		jerr(w, 502, "query: "+err.Error())
		return
	}
	pois, err := scanPois(rows)
	rows.Close()
	if err != nil {
		jerr(w, 502, err.Error())
		return
	}
	writePois(w, pois)
}

// point/radius parse the lat/lng (degrees) and r (metres) query params.
func point(q map[string][]string) (lat, lng float64, err error) {
	lat, err = strconv.ParseFloat(first(q, "lat"), 64)
	if err != nil {
		return 0, 0, fmt.Errorf("lat is required (decimal degrees)")
	}
	lng, err = strconv.ParseFloat(first(q, "lng"), 64)
	if err != nil {
		return 0, 0, fmt.Errorf("lng is required (decimal degrees)")
	}
	if lat < -90 || lat > 90 || lng < -180 || lng > 180 {
		return 0, 0, fmt.Errorf("lat/lng out of range")
	}
	return lat, lng, nil
}

func radius(q map[string][]string) (float64, error) {
	s := first(q, "r")
	if s == "" {
		return 500, nil
	}
	v, err := strconv.ParseFloat(s, 64)
	if err != nil || v < 50 || v > 50000 {
		return 0, fmt.Errorf("r must be 50..50000 metres")
	}
	return v, nil
}

// parsePoints: "lng,lat;lng,lat" → [[lng,lat],…] in that order.
func parsePoints(s string) [][2]float64 {
	var out [][2]float64
	for _, pair := range strings.Split(s, ";") {
		pair = strings.TrimSpace(pair)
		if pair == "" {
			continue
		}
		xy := strings.Split(pair, ",")
		if len(xy) != 2 {
			continue
		}
		lng, e1 := strconv.ParseFloat(strings.TrimSpace(xy[0]), 64)
		lat, e2 := strconv.ParseFloat(strings.TrimSpace(xy[1]), 64)
		if e1 == nil && e2 == nil {
			out = append(out, [2]float64{lng, lat})
		}
	}
	return out
}

// renumber resolves the placeholder scheme: the query body was written with
// $1..$4 for its own params (lat, lng, radius, limit) and the WHERE clause
// carries @@i@@ tokens for the n filter args, which come FIRST in the arg
// list. So: body $d → $(n+d), then @@i@@ → $i. n = number of filter args
// (= len(args) minus the body params, known at the call site).
func renumber(query string, n int) string {
	var b strings.Builder
	for i := 0; i < len(query); i++ {
		if query[i] == '$' && i+1 < len(query) && query[i+1] >= '1' && query[i+1] <= '4' {
			fmt.Fprintf(&b, "$%d", n+int(query[i+1]-'0'))
			i++
			continue
		}
		if strings.HasPrefix(query[i:], "@@") {
			end := strings.Index(query[i+2:], "@@")
			if end >= 0 {
				if d, err := strconv.Atoi(query[i+2 : i+2+end]); err == nil {
					fmt.Fprintf(&b, "$%d", d)
					i += 2 + end + 1 // loop's i++ finishes the last '@'
					continue
				}
			}
		}
		b.WriteByte(query[i])
	}
	return b.String()
}