using LibPQ, JSON, Graphs, DataFrames
config = JSON.parsefile("./appconfig.json")
host_url, _port = split(config["externalservice"]["sommpanion_db"]["url"], ':')
port = parse(Int, _port)
dbname = "winedb"
user = config["externalservice"]["sommpanion_db"]["user"]
password = config["externalservice"]["sommpanion_db"]["password"]
pg_conn_str = "host=$host_url port=$port dbname=$dbname user=$user password=$password"
db_connection = LibPQ.Connection(pg_conn_str)





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


using LibPQ
using DataFrames

"""
    extract_column_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_column_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_column_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."

# use only text content
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 - Distances.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 - Distances.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 find_related_tables_for_user_question(question::String)
  metadata_df = extract_column_metadata(pg_conn_str)
  embedding_ready = generate_embedding_payloads(metadata_df)

  # use only text content
  embedding_ready_2 = [i["text_content"] for i in embedding_ready]
  table_embedding = get_embedding(nats_conn, embedding_ready_2)
  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 - Distances.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

  #WORKING extract top 20 rows of sorted_df to get tables that related to user question
  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
  return optimized_context = resolve_semantic_cluster(vector_hits, g, table_to_id, id_to_table)
end