using LibPQ
using DataFrames

"""
    extract_vector_metadata(pg_conn_str::String) -> DataFrame

Queries PostgreSQL system catalogs to extract a rich semantic text map of every 
column in the database. Returns a DataFrame designed for vector embedding generation.
"""
function extract_vector_metadata(pg_conn_str::String)
    conn = LibPQ.Connection(pg_conn_str)
    
    # 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 = 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(df::DataFrame) -> Vector{Dict}

Transforms the metadata DataFrame into structured text strings optimal for 
vector space mapping.
"""
function generate_embedding_payloads(df::DataFrame)
    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(vector_hits::Vector{String}, g::SimpleGraph, table_to_id::Dict{String, Int}, id_to_table::Dict{Int, String}) -> Vector{String}

Takes a scattered array of semantically matched tables from Stage 1, navigates 
the undirected network structure, and isolates the minimum interconnected subgraph 
required to weave ALL hits into a single valid SQL query.
"""
function resolve_semantic_cluster(
    vector_hits::Vector{String}, 
    g::SimpleGraph, 
    table_to_id::Dict{String, Int}, 
    id_to_table::Dict{Int, 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

function get_embedding(nats_conn::NATS.Connection, text::AbstractArray{String})
  documents_dict = Dict("documents" => text)
  payloads = [("documents", documents_dict, "dictionary")]
  _, msg_envelope_json_str = msghandler.smartpack(
    config["externalservice"]["servicesloadbalancer"]["nats"], 
    payloads;   
    msg_purpose="embedding",  
    broker_url=config["nats_server_info"]["url"],
    fileserver_url=config["externalservice"]["fileserver"]["url"])

  reply = NATS.request(nats_conn, 
                      config["externalservice"]["servicesloadbalancer"]["nats"], 
                      msg_envelope_json_str, timeout=120)
  incoming_env_json_str = String(reply.payload)
  incoming_env = msghandler.smartunpack(incoming_env_json_str)
  embedding_response = incoming_env["payloads"][1][2]

  return embedding_response
end

nats_conn  = NATS.connect(config["nats_server_info"]["url"])

# Run the extractor
metadata_df = extract_vector_metadata(pg_conn_str)
embedding_ready = generate_embedding_payloads(metadata_df)

println(embedding_ready[1]["text_content"])
# Output: "Table: join_table | Column: seller_id | Type: integer [PRIMARY KEY] [FOREIGN KEY RELATIONAL LINK] | Description: Links unique sellers to their corresponding product items."

embedding_ready_2 = [i["text_content"] for i in embedding_ready]
table_embedding = get_embedding(nats_conn, embedding_ready_2)

user_question = 
  """
  Retrieves ["winery", "wine_name", "wine_id", "vintage", "region", "country", "wine_type", "grape", "serving_temperature", "sweetness", "intensity", "tannin", "acidity", "tasting_notes", "price", "currency", "image_url", "retailer_name", "retailer_id"] of wines that match the following criteria - {wine_name: Montrachet Grand Cru, winery: Domaine Jacques Prieur, region: Montrachet, country: France, , retailer_name: Yiem Wines Ltd, retailerid: f54eab6b-7650-4448-b009-c53f3efbcc3b}
  """
user_question_embedding = get_embedding(nats_conn, [user_question])


using Distances
similarity = 1 - cosine_dist(Float64.(table_embedding["data"][1]["embedding"]), 
                            Float64.(user_question_embedding["data"][1]["embedding"])
                            )
user_question_embedding = Float64.(user_question_embedding["data"][1]["embedding"])
user_question_similarity = []
for i in table_embedding["data"]
  i_data = i["embedding"]
  i_float = Float64.(i_data)
  r = 1 - cosine_dist(i_float, user_question_embedding)
  push!(user_question_similarity, r)
end

new_df = hcat(metadata_df, DataFrame(user_question_similarity = user_question_similarity))
sorted_df = sort(new_df, :user_question_similarity, rev=true) # sort max to min

# top 20 of sorted_df get this tables
vector_hits = ["retailer_wine", "wine", "wine_food", "retailer"]

g, id_to_table, table_to_id = harvest_undirected_schema_graph(pg_conn_str)

# tables that I should put schema in LLM context
optimized_context = resolve_semantic_cluster(vector_hits, g, table_to_id, id_to_table)



function related_tables_for_user_question()




end

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


# Agent 3 (The Entity Resolver): Instantly runs a fast, local token search (like BM25) to map messy user text (like HandOld) to the exact database string (Hand Old Bar & Grill) before the SQL is drafted.

using StringDistances

"""
    harvest_entity_catalog(conn_str::String, table::String, column::String) -> Vector{String}

Pulls unique, clean text strings from a specific entity column to build a local index.
"""
function harvest_entity_catalog(conn_str::String, table::String, column::String)
    conn = LibPQ.Connection(conn_str)
    
    # 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(execute(conn, query))
        # Return as a clean array of strings
        return String.(strip.(df[:, 1]))
    finally
        close(conn)
    end
end


"""
    resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.6) -> String

Parses user text, matches it against the real database catalog, and returns 
the exact string found in the database. Returns an empty string if no confident match.
"""
function resolve_entity(messy_input::String, catalog::Vector{String}; threshold=0.5)
    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



winery_catalog = harvest_entity_catalog(conn_str, "wine", "winery")
# Let's assume the catalog contains: ["Hand Old Bar & Grill", "Bangkok Diner", "Phuket Seafood"]

# 2. The user asks a messy question with a typo and compressed text
user_question = "What are the total sales at HandOld last week?"

# 3. Agent 3 isolates potential nouns or scans the question against the index
# We look for words that don't match standard english dictionary tokens, or check the full string segments
detected_entity = "Jacob" 

# 4. Run the resolution engine
exact_db_string = resolve_entity(detected_entity, winery_catalog)
# "United States"

println("Messy Input:   ", detected_entity)
println("Resolved Engine Value: ", exact_db_string)
# Output: "Hand Old Bar & Grill"