module llmUtil

export formatLLMtext, extractthink, checkAgentResponse_JSON, clean_json_response,
        extract_column_metadata, generate_embedding_payloads, resolve_semantic_cluster,
        harvest_entity_catalog, resolve_entity, harvest_db_undirected_schema_graph,
        get_db_table_schema

using UUIDs, JSON, Dates, DataFrames, StringDistances, Graphs, LibPQ
using ..util

# ---------------------------------------------- 100 --------------------------------------------- #

function formatLLMtext_qwen3(name::T, text::T;
  assistantStarter::Bool=false) where {T<:AbstractString}
  formattedtext = 
  if name == "system"
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  else
    """
    <|im_start|>$name
    $text
    <|im_end|>
    """
  end

  if assistantStarter
    formattedtext *=
      """
      <|im_start|>assistant
      """
  end

  return formattedtext
end


""" Convert a vector of chat message dictionaries into LLM model instruct format.

# Arguments
  - `messages::Vector{Dict{Symbol, T}}`
    A vector of dictionaries where each dictionary contains the keys `:name` (the name of the message owner) and `:text` (the text of the message).
  - `formatname::T`
    The name of the format to be used for converting the chat messages.
# Return
  - `formattedtext::String`
    text formatted to model format

# Example
```jldoctest
julia> using Revise
julia> using GeneralUtils
julia> chatmessage = [
          Dict(:name=> "system",:text=> "You are a helpful, respectful and honest assistant.",),
          Dict(:name=> "user",:text=> "list me all planets in our solar system.",),
          Dict(:name=> "assistant",:text=> "I'm sorry. I don't know. You tell me.",),
        ]
julia> formattedtext = GeneralUtils.formatLLMtext(chatmessage, "llama3instruct")
"<|begin_of_text|>\n  <|start_header_id|>system<|end_header_id|>\n  You are a helpful, respectful and honest assistant.\n  <|eot_id|>\n  <|start_header_id|>user<|end_header_id|>\n  list me all planets in our solar system.\n  <|eot_id|>\n  <|start_header_id|>assistant<|end_header_id|>\n  I'm sorry. I don't know. You tell me.\n  <|eot_id|>\n"
```
"""
function formatLLMtext(messages::Vector{Dict{Symbol, T}}, formatname::String
  )::String where {T<:AbstractString}
  f = 
    if formatname == "llama3instruct"
      formatLLMtext_llama3instruct
    elseif formatname == "mistral"
      # not define yet
    elseif formatname == "phi3instruct"
      # not define yet
    elseif formatname == "qwen"
      formatLLMtext_qwen
    elseif formatname == "qwen3"
      formatLLMtext_qwen3
    elseif formatname == "phi4"
      formatLLMtext_phi4
    elseif formatname == "granite3"
      formatLLMtext_granite3
    else
      error("$formatname template not define yet")
    end

  str = ""
  for (i, t) in enumerate(messages)
    if i < length(messages)
      str *= f(t[:name], t[:text])
    else
      str *= f(t[:name], t[:text]; assistantStarter=true)
    end
  end

  return str
end


function extractthink(text::String)
  think = nothing
  str = nothing
  if occursin("<think>", text)
    r = GeneralUtils.extractTextBetweenString(text, "<think>", "</think>")
    if r[:success]
      think = r[:text]
    end
    str = string(split(text, "</think>")[2])
  else
    str = text
  end
  return think, str
end



""" Validate that an agent's JSON response contains all required keys and no extra keys.

The function checks if `responsedict` contains exactly the keys specified in `requiredKeys`
with no duplicates and no missing keys. It is designed to validate structured agent responses
against an expected schema.

# Arguments
- `responsedict::Dict`
  A dictionary containing the agent's JSON response. Must be a plain `Dict` or
  dictionary-like object with String keys.
- `requiredKeys::T` where `T<:Array{String}`
  An array of required key names that must be present in `responsedict`.

# Return
- `Tuple{Bool, Union{String, Nothing}}`
  A tuple where the first element indicates whether validation passed (`true`) or failed (`false`),
  and the second element contains an error message if validation failed, or `nothing` if it passed.

# Details
The validation logic checks:
1. **Duplicate keys**: If `responsedict` contains more keys than `requiredKeys`, validation fails
   because the agent included extra/unexpected keys.
2. **Missing keys**: If any key in `requiredKeys` is absent from `responsedict`, validation fails
   and the specific missing keys are listed in the error message.
3. **Valid response**: If all required keys are present and no extra keys exist, validation passes.

# Example

```julia
julia> using GeneralUtils
julia> requiredKeys = ["wine_name", "price", "rating"]
julia> response = Dict("wine_name"=>"Château Margaux", "price"=>250.0, "rating"=>98)
julia> GeneralUtils.checkAgentResponse_JSON(response, requiredKeys)
(true, nothing)

julia> response_missing = Dict("wine_name"=>"Château Margaux", "price"=>250.0)
julia> GeneralUtils.checkAgentResponse_JSON(response_missing, requiredKeys)
(false, "rating are missing from your previous response")

julia> response_extra = Dict("wine_name"=>"Château Margaux", "price"=>250.0, "rating"=>98, "extra_field"=>"data")
julia> GeneralUtils.checkAgentResponse_JSON(response_extra, requiredKeys)
(false, "Your previous attempt has duplicated points according to the required response format")
```
"""
function checkAgentResponse_JSON(responsedict::T1, requiredKeys::T2
  )::Tuple where {T1<:AbstractDict, T2<:Array{String}}
  _responsedictKey = keys(responsedict)
  responsedictKey = [i for i in _responsedictKey] # convert into a list
  is_requiredKeys_in_responsedictKey = [i ∈ responsedictKey for i in requiredKeys]
  ispass = false
  errormsg = nothing
  if length(is_requiredKeys_in_responsedictKey) > length(requiredKeys)
    errormsg = "Your previous attempt has duplicated points according to the required response format"
    ispass = false
  elseif !all(is_requiredKeys_in_responsedictKey)
    zeroind = findall(x -> x == 0, is_requiredKeys_in_responsedictKey)
    missingkeys = [requiredKeys[i] for i in zeroind]
    errormsg = "$missingkeys are missing from your previous response"
    ispass = false
  else
    ispass = true
  end
  return (ispass, errormsg)
end


""" Convert a plain text string containing key-value pairs into a JSON-formatted string.

This function takes text containing key-value pairs (typically extracted from LLM responses)
and wraps them in proper JSON braces to create a valid JSON string. It cleans the input
by removing common formatting artifacts like braces, code block markers, and language
specifiers before wrapping the content.

# Arguments
- `text::String`
  A string containing key-value pairs, typically in the format `key: value, key: value`.
  May contain leading/trailing braces, code block markers (```), or language specifiers
  (e.g., `json`) that will be removed.

# Return
- `String`
  A JSON-formatted string with the key-value pairs wrapped in `{}` braces.

# Notes
- The function removes `{`, `}`, `````, and `json` from the input before wrapping.
- The output is always wrapped in curly braces to create valid JSON structure.
- This is typically used to sanitize LLM responses that contain key-value data
  but may include formatting artifacts.

# Examples
```jldoctest
julia> text = "thought: Hello, action: CHATBOX, input: How can I help?"
julia> clean_json_response(text)
"{thought: Hello, action: CHATBOX, input: How can I help?}"

julia> text = "{thought: Hello, action: CHATBOX}"
julia> clean_json_response(text)
"{thought: Hello, action: CHATBOX}"

julia> text = "```json{thought: Hello, action: CHATBOX}```"
julia> clean_json_response(text)
"{thought: Hello, action: CHATBOX}"
```
"""
function clean_json_response(text::String)
  removelist = ["{", "}", "```", "json", "\n"]
  return '{' * removestring(text, removelist) * '}'
end



""" Harvest database schema as undirected graph from PostgreSQL.

Extracts table structure and relationships from a PostgreSQL database by querying
system catalogs to build a graph representation of tables and their relationships.

# Arguments
- `pg_conn_str::String`
  PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")

# Return
- `g::SimpleGraph`: An undirected graph where nodes represent tables and edges represent
  relationships (foreign keys or shared ID column patterns).
- `id_to_table::Dict{Int, String}`: Maps numerical node IDs (1..n) to actual table names.
- `table_to_id::Dict{String, Int}`: Reverse mapping from table names to graph node IDs.

# Details
The function extracts schema information using two strategies:
1. **Explicit Foreign Keys**: Queries `pg_constraint` for actual foreign key relationships
2. **Fallback Name Matching**: Infers relationships from shared column naming patterns
   (e.g., `seller_id`, `product_id` columns across tables)

# Example
```julia
julia> using GeneralUtils
julia> pg_conn_str = "host=localhost port=5432 dbname=winedb user=admin password=secret"
julia> g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
julia> vertices(g)
10
```

# Integration Guide: Finding Related Tables for User Questions
This function is designed to work together with `extract_column_metadata` and `resolve_semantic_cluster` to answer user questions by identifying related tables:

```julia
# Step 1: Extract column metadata and generate embeddings for semantic search
pg_conn_str = "host=localhost port=5432 dbname=winedb user=admin password=secret"
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)

# Use only text content for embedding
embedding_ready_2 = [i["text_content"] for i in embedding_ready]
table_embedding = get_embedding(embedding_ready_2)

# Get embedding for user question
_user_question_embedding = get_embedding([question])
user_question_embedding = Float64.(_user_question_embedding["data"][1]["embedding"])

# Calculate similarity between question and all columns
user_question_similarity = []
for i in table_embedding["data"]
  i_data = i["embedding"]
  i_float = Float64.(i_data)
  r = 1 - Distances.cosine_dist(i_float, user_question_embedding)
  push!(user_question_similarity, r)
end

# Step 2: Find top related tables
new_df = hcat(metadata_df, DataFrame(user_question_similarity = user_question_similarity))
sorted_df = sort(new_df, :user_question_similarity, rev=true)
_top_20_tables = unique(sorted_df[1:20, :table_name])
top_20_tables = [i for i in _top_20_tables]

# Step 3: Build schema graph and resolve table relationships
g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
table_relationship = GeneralUtils.resolve_semantic_cluster(top_20_tables, g, table_to_id, id_to_table)

# table_relationship now contains tables in the order they should be joined
```
"""
function harvest_db_undirected_schema_graph(pg_conn_str)   
  conn = LibPQ.Connection(pg_conn_str) 
    try
      # 1. Fetch all user tables
      table_query = """
          SELECT c.relname AS table_name, c.oid
          FROM pg_class c
          JOIN pg_namespace n ON n.oid = c.relnamespace
          WHERE n.nspname = 'public' AND c.relkind = 'r';
      """
      table_df = DataFrame(execute(conn, table_query))
      table_names = table_df.table_name
      
      num_tables = length(table_names)
      table_to_id = Dict{String, Int}(name => i for (i, name) in enumerate(table_names))
      id_to_table = Dict{Int, String}(i => name for (i, name) in enumerate(table_names))
      
      # CRITICAL: Use SimpleGraph (Undirected) so pathfinding can traverse both ways
      g = SimpleGraph(num_tables)
      
      # 2. Strategy A: Extract Explicit Foreign Keys
      fk_query = """
          SELECT 
              conrelid::regclass::text AS source_table,
              confrelid::regclass::text AS target_table
          FROM pg_constraint c
          JOIN pg_namespace n ON n.oid = c.connamespace
          WHERE c.contype = 'f' AND n.nspname = 'public';
      """
      fk_df = DataFrame(execute(conn, fk_query))
      
      for row in eachrow(fk_df)
          src = split(replace(row.source_table, "\"" => ""), '.')[end]
          tgt = split(replace(row.target_table, "\"" => ""), '.')[end]
          
          if haskey(table_to_id, src) && haskey(table_to_id, tgt)
              add_edge!(g, table_to_id[src], table_to_id[tgt])
          end
      end
      
      # 3. Strategy B: Fallback Name Matching (For missing explicit constraints)
      # Fetch all columns for all tables
      col_query = """
          SELECT a.attname AS column_name, c.relname AS table_name
          FROM pg_attribute a
          JOIN pg_class c ON c.oid = a.attrelid
          JOIN pg_namespace n ON n.oid = c.relnamespace
          WHERE n.nspname = 'public' AND c.relkind = 'r' AND a.attnum > 0 AND NOT a.attisdropped;
      """
      col_df = DataFrame(execute(conn, col_query))
      
      # Group columns by name to see who shares what fields
      for gdf in groupby(col_df, :column_name)
        col_name = gdf.column_name[1]
        
        # We only infer connections on ID fields (e.g., seller_id, product_id)
        # Avoid matching generic names like 'id', 'created_at', or 'name'
        if endswith(lowercase(col_name), "_id") && col_name != "id"
          sharing_tables = gdf.table_name
          
          # Connect all tables that share this ID column
          for i in eachindex(sharing_tables), j in (i+1):length(sharing_tables)
            t1 = sharing_tables[i]
            t2 = sharing_tables[j]
            
            if haskey(table_to_id, t1) && haskey(table_to_id, t2)
              add_edge!(g, table_to_id[t1], table_to_id[t2])
            end
          end
        end
      end
      
      return g, id_to_table, table_to_id

    finally
      close(conn)
    end
end



""" Extract vector metadata from PostgreSQL database.

Queries PostgreSQL system catalogs to extract column metadata including table names,
column names, data types, descriptions, and constraint information (primary/foreign keys).
Returns a DataFrame designed for vector embedding generation.

# Arguments
- `pg_conn_str::String`
  PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")

# Return
- `DataFrame`
  A DataFrame with columns:
  - `table_name`: Name of the table
  - `column_name`: Name of the column
  - `data_type`: PostgreSQL data type with typemod
  - `column_description`: Column's comment/description (empty string if none)
  - `is_primary_key`: Boolean indicating if column is part of primary key
  - `is_foreign_key`: Boolean indicating if column is part of foreign key
  - `vector_id`: Generated unique identifier in format `col_{table_name}_{column_name}`

# Example
```julia
julia> using GeneralUtils
julia> pg_conn = "host=localhost port=5432 dbname=winedb user=admin password=secret"
julia> df = GeneralUtils.extract_column_metadata(pg_conn)
DataFrame
6 rows × 7 columns
table_name  column_name  data_type  column_description  is_primary_key  is_foreign_key  vector_id
─────────────┬───────────┬───────────┬───────────────────┬───────────────┬───────────────┬────────────────────────
users       id           integer    User ID             true            false         col_users_id
users       name         text       User name           false           false         col_users_name
users       email        text       User email          false           false         col_users_email
products    id           integer    Product ID          true            false         col_products_id
products    price        numeric    Product price       false           false         col_products_price
products    user_id      integer    Reference to user   false           true          col_products_user_id
```

# Integration Guide
This function is the first step in the semantic table discovery pipeline. Use it with `generate_embedding_payloads` to prepare data for vector similarity search:

```julia
# Step 1: Extract metadata
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)

# Step 2: Generate embedding payloads for semantic search
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)

# Use text_content field for embedding
text_content_list = [i["text_content"] for i in embedding_ready]
embeddings = get_embedding(text_content_list)
```
"""
function extract_column_metadata(pg_conn_str::String)::DataFrame
  conn = LibPQ.Connection(pg_conn_str)
  return extract_column_metadata(conn)
end

function extract_column_metadata(conn::LibPQ.Connection)::DataFrame
  # This direct SQL query pulls the column specifications along with column-level descriptions
  query = """
        SELECT 
            c.relname AS table_name,
            a.attname AS column_name,
            format_type(a.atttypid, a.atttypmod) AS data_type,
            COALESCE(d.description, '') AS column_description,
            CASE WHEN pk.contype = 'p' THEN true ELSE false END AS is_primary_key,
            CASE WHEN fk.contype = 'f' THEN true ELSE false END AS is_foreign_key
        FROM pg_attribute a
        JOIN pg_class c ON c.oid = a.attrelid
        JOIN pg_namespace n ON n.oid = c.relnamespace
        -- Join to fetch column comments/descriptions
        LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
        -- Check if column is part of a Primary Key
        LEFT JOIN pg_constraint pk ON pk.conrelid = c.oid 
            AND pk.contype = 'p' 
            AND a.attnum = ANY(pk.conkey)
        -- Check if column is part of a Foreign Key
        LEFT JOIN pg_constraint fk ON fk.conrelid = c.oid 
            AND fk.contype = 'f' 
            AND a.attnum = ANY(fk.conkey)
        WHERE 
            n.nspname = 'public'          -- Only user schemas
            AND c.relkind = 'r'            -- Only standard tables
            AND a.attnum > 0               -- Skip system hidden columns
            AND NOT a.attisdropped;        -- Skip dropped columns
    """
  
  try
      # Execute and format into a clean DataFrame
      result = LibPQ.execute(conn, query)
      df = DataFrame(result)
      
      # Create a unique document ID for each vector row
      df.vector_id = ["col_\$(row.table_name)_\$(row.column_name)" for row in eachrow(df)]
      
      return df
  finally
      close(conn)
  end
end


""" Generate embedding payloads from vector metadata DataFrame.

Transforms the metadata DataFrame from `extract_column_metadata` into a vector of
dictionaries structured for vector embedding storage and retrieval.

# Arguments
- `df::DataFrame`
  A DataFrame with columns from `extract_column_metadata`: `table_name`, `column_name`,
  `data_type`, `column_description`, `is_primary_key`, `is_foreign_key`, `vector_id`.

# Return
- `Vector{Dict}`
  A vector of dictionaries with keys:
  - `id`: The vector ID from `vector_id`
  - `text_content`: A structured text string combining table, column, type, constraint
    markers, and description
  - `metadata`: A dictionary containing `table`, `column`, and `type`

# Details
The function constructs rich text payloads by:
1. Building a base string with table name, column name, and data type
2. Appending constraint markers `[PRIMARY KEY]` or `[FOREIGN KEY RELATIONAL LINK]`
3. Adding column description if available, otherwise generating a default description

# Example
```julia
julia> using GeneralUtils
julia> pg_conn = "host=localhost port=5432 dbname=winedb user=admin password=secret"
julia> df = GeneralUtils.extract_column_metadata(pg_conn)
julia> payloads = GeneralUtils.generate_embedding_payloads(df)
3-element Vector{Dict}:
 Dict("id" => "col_users_id", "text_content" => "Table: users | Column: id | Type: integer [PRIMARY KEY] | Description: User ID", "metadata" => Dict("table" => "users", "column" => "id", "type" => "integer"))
 Dict("id" => "col_users_name", "text_content" => "Table: users | Column: name | Type: text | Description: User name", "metadata" => Dict("table" => "users", "column" => "name", "type" => "text"))
 Dict("id" => "col_users_email", "text_content" => "Table: users | Column: email | Type: text | Description: User email", "metadata" => Dict("table" => "users", "column" => "email", "type" => "text"))
```

# Integration Guide
This function transforms column metadata into a format suitable for semantic search.
Use `text_content` field to generate embeddings, then compare against user question embeddings:

```julia
# 1. Extract metadata
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)

# 2. Generate payloads
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)

# 3. Extract text content for embedding (only this field is used for similarity)
text_content_list = [i["text_content"] for i in embedding_ready]

# 4. Generate embeddings for all columns
column_embeddings = get_embedding(text_content_list)

# 5. Generate embedding for user question
question_embedding = get_embedding([user_question])

# 6. Calculate cosine similarity to find most relevant columns/tables
# (See harvest_db_undirected_schema_graph for full integration example)
```
"""
function generate_embedding_payloads(df::DataFrame)::Vector{Dict}
    payloads = Dict[]
    
    for row in eachrow(df)
        # 1. Build a rich text description summarizing the column's role
        text_payload = "Table: $(row.table_name) | Column: $(row.column_name) | Type: $(row.data_type)"
        
        if row.is_primary_key
            text_payload *= " [PRIMARY KEY]"
        end
        if row.is_foreign_key
            text_payload *= " [FOREIGN KEY RELATIONAL LINK]"
        end
        
        # Append business descriptions if they exist in the database comments
        if !isempty(strip(row.column_description))
            text_payload *= " | Description: $(row.column_description)"
        else
            text_payload *= " | Description: Represents $(row.column_name) data fields within the $(row.table_name) architecture."
        end
        
        # 2. Package everything neatly to be passed to your vector store client
        push!(payloads, Dict(
            "id" => row.vector_id,
            "text_content" => text_payload,
            "metadata" => Dict(
                "table" => row.table_name,
                "column" => row.column_name,
                "type" => row.data_type
            )
        ))
    end
    
    return payloads
end


""" Resolve semantic cluster from vector hits using graph traversal.

Finds the minimum interconnected subgraph that connects all semantically matched tables
from vector search results, enabling construction of valid SQL queries across related tables.

# Arguments
- `vector_hits::Vector{String}`
  A vector of table names matched semantically from Stage 1 vector search.
- `g::SimpleGraph`
  An undirected graph representing table relationships (nodes=tables, edges=foreign key relations).
- `table_to_id::Dict{String, Int}`
  Mapping from table names to node IDs in the graph.
- `id_to_table::Dict{Int, String}`
  Reverse mapping from node IDs to table names.

# Return
- `Vector{String}`
  A vector of table names representing the minimum subgraph that connects all input vector hits.
  The order reflects traversal path from anchor node to connected components.

# Details
The algorithm:
1. Validates vector hits against the graph's table mapping
2. Handles edge cases: empty hits, single table (returns as-is)
3. Uses the highest-ranked vector hit as anchor node
4. For each remaining hit, finds shortest path to current subgraph using Dijkstra's algorithm
5. Builds minimal connected subgraph containing all hits
6. Returns table names in traversal order

# Example
```julia
julia> using GeneralUtils, Graphs
julia> g = SimpleGraph(5)
julia> add_edge!(g, 1, 2)
julia> add_edge!(g, 2, 3)
julia> add_edge!(g, 3, 4)
julia> table_to_id = Dict("users" => 1, "orders" => 2, "payments" => 3, "products" => 4, "inventory" => 5)
julia> id_to_table = Dict(1 => "users", 2 => "orders", 3 => "payments", 4 => "products", 5 => "inventory")
julia> vector_hits = ["users", "products"]
julia> GeneralUtils.resolve_semantic_cluster(vector_hits, g, table_to_id, id_to_table)
["users", "orders", "products"]
```

# Integration Guide: Understanding Table Relationships
This function determines how tables from vector search results are connected in the database schema.
The returned list represents the optimal join order for constructing SQL queries.

**How tables are linked:**
- **Explicit Foreign Keys**: Direct relationships defined by `FOREIGN KEY` constraints in PostgreSQL
- **Implicit Relationships**: Tables sharing similar column naming patterns (e.g., `user_id` in both `users` and `orders` tables)

**Example output interpretation:**
```julia
# Input: Top 20 tables from semantic search
top_tables = ["users", "products", "payments"]

# Output: Tables in join order
table_relationship = ["users", "orders", "payments"]
# Interpretation: 
#   1. Start with 'users' table
#   2. Join 'orders' via foreign key (likely users.id -> orders.user_id)
#   3. Join 'payments' via foreign key (likely orders.id -> payments.order_id)
```

**Full workflow example:**
```julia
# Step 1: Get column embeddings and find semantically related tables
metadata_df = GeneralUtils.extract_column_metadata(pg_conn_str)
embedding_ready = GeneralUtils.generate_embedding_payloads(metadata_df)
text_content = [i["text_content"] for i in embedding_ready]
table_embedding = get_embedding(text_content)

# Step 2: Calculate similarity with user question
question_embedding = Float64.(get_embedding([question])["data"][1]["embedding"])
similarities = [1 - Distances.cosine_dist(Float64.(i["embedding"]), question_embedding) 
                for i in table_embedding["data"]]

# Step 3: Extract top related tables
top_tables = unique(metadata_df[sortperm(similarities)[1:20], :table_name])

# Step 4: Build schema graph and resolve relationships
g, id_to_table, table_to_id = GeneralUtils.harvest_db_undirected_schema_graph(pg_conn_str)
related_tables = GeneralUtils.resolve_semantic_cluster(top_tables, g, table_to_id, id_to_table)

# related_tables now contains tables in optimal join order for SQL query construction
```
"""
function resolve_semantic_cluster(
    vector_hits::Vector{String}, 
    g::SimpleGraph, 
    table_to_id::Dict{String, Int}, 
    id_to_table::Dict{Int, String}
  )::Vector{String}
    # Filter out hits that don't exist in our actual database graph mapping
    valid_node_ids = Int[]
    for hit in vector_hits
        if haskey(table_to_id, hit)
            push!(valid_node_ids, table_to_id[hit])
        else
            @warn "Vector hit '$hit' does not map to an existing database table."
        end
    end
    
    unique!(valid_node_ids)
    
    # Edge Case Handlers
    if isempty(valid_node_ids)
        return String[]
    elseif length(valid_node_ids) == 1
        return [id_to_table[valid_node_ids[1]]]
    end

    # The Isolated Subgraph Set to build our final context
    schema_subgraph_nodes = Set{Int}()
    
    # Phase A: Select an initial anchor component. We use the highest-ranked vector hit.
    anchor_node = valid_node_ids[1]
    push!(schema_subgraph_nodes, anchor_node)
    
    # Phase B: Sequentially route paths to all other semantic coordinates
    for target_node in valid_node_ids[2:end]
        # Skip if an earlier loop trajectory already naturally absorbed this table
        if target_node in schema_subgraph_nodes
            continue
        end
        
        # Calculate the shortest path tree from the CURRENT state of our subgraph
        # We find the shortest path from the target back to ANY node currently in our tree
        shortest_paths = dijkstra_shortest_paths(g, target_node)
        
        # Find which node currently in our subgraph is closest to the target node
        closest_subgraph_node = 0
        min_distance = Inf
        
        for subgraph_node in schema_subgraph_nodes
            dist = shortest_paths.dists[subgraph_node]
            if dist < min_distance
                min_distance = dist
                closest_subgraph_node = subgraph_node
            end
        end
        
        # Reconstruct the path from the target node to the closest point on our existing tree
        if closest_subgraph_node != 0
            curr = closest_subgraph_node
            while curr != 0
                push!(schema_subgraph_nodes, curr)
                curr = shortest_paths.parents[curr]
                if curr == target_node
                    push!(schema_subgraph_nodes, target_node)
                    break
                end
            end
        end
    end
    
    # Map the unique structural nodes back to clean table names
    return [id_to_table[node_id] for node_id in schema_subgraph_nodes]
end


""" Harvest entity catalog from database column.

Extracts unique, non-null values from a specific column to build a local index for 
semantic search or entity resolution.

# Arguments
- `pg_conn_str::String`
  PostgreSQL connection string in LibPQ format (e.g., "host=hostname port=5432 dbname=database user=username password=secret")
- `table::String`
  Table name to query
- `column::String`
  Column name containing entity values

# Return
- `Vector{String}`
  A vector of unique, stripped strings from the specified column. Empty strings
  are removed via `strip()`.

# Details
The function:
1. Connects to PostgreSQL database
2. Executes `SELECT DISTINCT column FROM table WHERE column IS NOT NULL`
3. Converts result to DataFrame
4. Strips whitespace from each value and converts to String
5. Returns clean vector of unique entity values

# Example
```julia
julia> using GeneralUtils
julia> conn = "host=localhost port=5432 dbname=winedb user=admin password=secret"
julia> fruits = GeneralUtils.harvest_entity_catalog(conn, "products", "fruit_name")
["Apple", "Banana", "Orange", "Mango"]
```

# Integration Guide: Finding Closest Entity Match in Database
When users type queries, they often make typos or use abbreviations. This function finds
the closest matching entity from a database column, handling common input errors.

**Example:**
```julia
# Database contains: ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
# User types: "HandOld" (compressed words, missing space)

catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "venues", "name")
closest = GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.9)
# Returns: "Hand Old Bar & Grill" (closest match)
```

**Step 1: Extract all valid values from a column**
```julia
# Get all wine names from the database
wine_names = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
# Returns: ["Château Margaux", "Château Lafite", "Domaine Leroy", ...]
```

**Step 2: Find closest match for user input**
```julia
# User types a wine name with a typo
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
closest_wine = GeneralUtils.resolve_entity("Chateu Margaux", catalog; threshold=0.9)
# Returns: "Château Margaux" (closest match)
```

**Step 3: Resolve all fields in a response**
```julia
# Agent produces a response with potentially misspelled values
responsedict = Dict("wine_name" => "Chateu Margaux", "region" => "Bord")

# For each field, find the closest match in the database
for (k, v) in responsedict
  catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", k)
  resolved = GeneralUtils.resolve_entity(v, catalog; threshold=0.9)
  responsedict[k] = resolved  # Replace messy input with database entity
end
```

**Threshold guidance:**
- `threshold=0.5` - Lenient, may return false positives
- `threshold=0.7` - Moderate balance
- `threshold=0.9` - Strict, only high-confidence matches
- Returns `""` when no match meets threshold
"""
function harvest_entity_catalog(pg_conn_str::String, table::String, column::String)::Vector{String}
  conn = LibPQ.Connection(pg_conn_str)
  return harvest_entity_catalog(conn, table, column)
end

function harvest_entity_catalog(conn::LibPQ.Connection, table::String, column::String)::Vector{String}
  
  # We only care about unique, non-null values to keep the index fast and dense
  query = "SELECT DISTINCT $(column) FROM $(table) WHERE $(column) IS NOT NULL;"
  
  try
    df = DataFrame(LibPQ.execute(conn, query))
    # Return as a clean array of strings
    return String.(strip.(df[:, 1]))
  catch
    return String[]
  finally
    close(conn)
  end
end


""" Resolve entity name from messy input using fuzzy string matching.

Matches user-provided text against a reference catalog using Jaro-Winkler similarity
and returns the closest matching exact string from the database catalog.

# Arguments
- `messy_input::String`
  The user input text that may contain typos, compressed words, or variations.
- `catalog::Vector{String}`
  A vector of valid, exact entity strings from the database.

# Keyword Arguments
- `threshold::Float64` (default: `0.5`)
  Minimum similarity score (0.0 to 1.0) required to return a match. Lower values
  allow more lenient matching; higher values require closer matches.

# Return
- `String`
  The exact matching string from `catalog` if similarity score ≥ threshold,
  otherwise an empty string `""`.

# Details
The function:
1. Normalizes input to lowercase and strips whitespace
2. Computes Jaro-Winkler similarity score against each catalog entry
3. Applies substring fallback: if compressed words match (e.g., "HandOld" → "Hand Old Bar & Grill"),
   boosts score to 0.85
4. Returns the highest-scoring catalog entry if score ≥ threshold, else empty string

# Example
```julia
julia> using GeneralUtils
julia> catalog = ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
julia> GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.5)
"Hand Old Bar & Grill"

julia> GeneralUtils.resolve_entity("Wine Cellar", catalog; threshold=0.5)
"Wine Cellar"

julia> GeneralUtils.resolve_entity("Unknown Place", catalog; threshold=0.5)
""
```

# Integration Guide: Finding Closest Entity Match in Database
When users type queries, they often make typos or use abbreviations. This function finds
the closest matching entity from a database column, handling common input errors.

**Example:**
```julia
# Database contains: ["Hand Old Bar & Grill", "Hand Old", "Wine Cellar"]
# User types: "HandOld" (compressed words, missing space)

catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "venues", "name")
closest = GeneralUtils.resolve_entity("HandOld", catalog; threshold=0.9)
# Returns: "Hand Old Bar & Grill" (closest match)
```

**Step 1: Extract all valid values from a column**
```julia
# Get all wine names from the database
wine_names = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
# Returns: ["Château Margaux", "Château Lafite", "Domaine Leroy", ...]
```

**Step 2: Find closest match for user input**
```julia
# User types a wine name with a typo
catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", "wine_name")
closest_wine = GeneralUtils.resolve_entity("Chateu Margaux", catalog; threshold=0.9)
# Returns: "Château Margaux" (closest match)
```

**Step 3: Resolve all fields in a response**
```julia
# Agent produces a response with potentially misspelled values
responsedict = Dict("wine_name" => "Chateu Margaux", "region" => "Bord")

# For each field, find the closest match in the database
for (k, v) in responsedict
  catalog = GeneralUtils.harvest_entity_catalog(pg_conn_str, "wine", k)
  resolved = GeneralUtils.resolve_entity(v, catalog; threshold=0.9)
  responsedict[k] = resolved  # Replace messy input with database entity
end
```

**Threshold guidance:**
- `threshold=0.5` - Lenient, may return false positives
- `threshold=0.7` - Moderate balance
- `threshold=0.9` - Strict, only high-confidence matches
- Returns `""` when no match meets threshold
"""
function resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.5)::String
    best_match = ""
    highest_score = 0.0
    
    # Normalize input text to ensure case-insensitive matching
    clean_input = lowercase(strip(messy_input))
    
    for real_string in catalog
        clean_real = lowercase(real_string)
        
        # Calculate phonetic/structural similarity score (0.0 to 1.0)
        # JaroWinkler is optimized for short strings, names, and partial acronyms
        score = compare(clean_real, clean_input, JaroWinkler())
        
        # Substring/Token fallback: handle cases like "HandOld" matching "Hand Old Bar & Grill"
        # We strip spaces to check if the user just compressed words together
        if contains(replace(clean_real, " " => ""), clean_input)
            score = max(score, 0.85) 
        end
        
        if score > highest_score
            highest_score = score
            best_match = real_string
        end
    end
    
    # Only return if we cross our safety confidence barrier
    if highest_score >= threshold
        return best_match
    end
    
    return "" # No confident match found
end


function get_db_table_schema(pg_conn_str::String, table_name::String)::DataFrame
  conn = LibPQ.Connection(pg_conn_str)
  return get_db_table_schema(conn, table_name)
end

function get_db_table_schema(conn::LibPQ.Connection, table_name::String)::DataFrame
  # This direct SQL query pulls the column specifications along with column-level descriptions
  query = """
        SELECT 
            a.attname AS column_name,
            format_type(a.atttypid, a.atttypmod) AS data_type,
            COALESCE(d.description, '') AS column_comment,
            CASE 
                WHEN p.contype = 'p' THEN 'PRIMARY KEY'
                WHEN p.contype = 'u' THEN 'UNIQUE'
                WHEN p.contype = 'f' THEN 'FOREIGN KEY'
                WHEN p.contype = 'c' THEN 'CHECK'
                ELSE ''
            END AS constraint_type,
            COALESCE(p.conname, '') AS constraint_name
        FROM 
            pg_catalog.pg_attribute a
        JOIN 
            pg_catalog.pg_class c ON a.attrelid = c.oid
        JOIN 
            pg_catalog.pg_namespace n ON c.relnamespace = n.oid
        LEFT JOIN 
            pg_catalog.pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
        LEFT JOIN 
            pg_catalog.pg_constraint p ON p.conrelid = c.oid AND a.attnum = ANY(p.conkey)
        WHERE 
            c.relname = '$table_name' -- <-- Put your table name here
            AND n.nspname = 'public'      -- <-- Change schema if not 'public'
            AND a.attnum > 0 
            AND NOT a.attisdropped
        ORDER BY 
            a.attnum;
    """
  
  try
      # Execute and format into a clean DataFrame
      result = LibPQ.execute(conn, query)
      df = DataFrame(result)
      
      return df
  finally
      close(conn)
  end
end















end # module llmUtil