# ============================================================
# SPID — knowledge graph builder
#
# Derives nodes (Company / Project / Technology / Location /
# Partner) and edges (owns / uses / located_in / partners_with)
# from the resolved projects, then exports to parquet + Cypher.
# ============================================================
from __future__ import annotations

from pathlib import Path

from spid.config import get_logger
from spid.store import Store

log = get_logger("graph")

_BUILD = """
DELETE FROM kg_nodes;
DELETE FROM kg_edges;

-- Company nodes (one per cik; a cik may carry varying name/sector across projects)
INSERT INTO kg_nodes
SELECT 'C:' || cik, 'Company', any_value(coalesce(company_name, ticker, cik)),
       any_value(to_json({'cik': cik, 'ticker': ticker, 'sector': sector}))
FROM projects WHERE cik IS NOT NULL
GROUP BY cik;

-- Project nodes
INSERT INTO kg_nodes
SELECT 'P:' || project_id, 'Project', project_name,
       to_json({'type': project_type, 'status': status,
                'amount_usd': total_amount_usd, 'n_filings': n_filings,
                'first_seen': first_seen, 'last_seen': last_seen})
FROM projects;

-- Technology nodes
INSERT INTO kg_nodes
SELECT DISTINCT 'T:' || t, 'Technology', t, NULL
FROM (SELECT unnest(technologies) AS t FROM projects) WHERE t IS NOT NULL AND t <> '';

-- Location nodes
INSERT INTO kg_nodes
SELECT DISTINCT 'L:' || canonical_location, 'Location', canonical_location, NULL
FROM projects WHERE canonical_location IS NOT NULL AND canonical_location <> '';

-- Partner nodes (from mentions; id is case-folded so collapse case variants to one row)
INSERT INTO kg_nodes
SELECT 'PR:' || lower(p), 'Partner', any_value(p), NULL
FROM (SELECT unnest(partners) AS p FROM project_mentions)
WHERE p IS NOT NULL AND p <> ''
GROUP BY lower(p);

-- Edges: Company -owns-> Project
INSERT INTO kg_edges
SELECT 'C:' || cik, 'owns', 'P:' || project_id,
       to_json({'amount_usd': total_amount_usd})
FROM projects WHERE cik IS NOT NULL;

-- Edges: Project -uses-> Technology
INSERT INTO kg_edges
SELECT 'P:' || project_id, 'uses', 'T:' || t, NULL
FROM (SELECT project_id, unnest(technologies) AS t FROM projects)
WHERE t IS NOT NULL AND t <> '';

-- Edges: Project -located_in-> Location
INSERT INTO kg_edges
SELECT 'P:' || project_id, 'located_in', 'L:' || canonical_location, NULL
FROM projects WHERE canonical_location IS NOT NULL AND canonical_location <> '';

-- Edges: Project -requires_capital (self-annotation kept as company edge prop already)
"""


def build_graph(store: Store) -> dict:
    store.con.execute(_BUILD)
    n_nodes = store.count("kg_nodes")
    n_edges = store.count("kg_edges")
    by_type = dict(store.con.execute(
        "SELECT node_type, COUNT(*) FROM kg_nodes GROUP BY 1 ORDER BY 2 DESC").fetchall())
    log.info("Knowledge graph: %d nodes, %d edges | %s", n_nodes, n_edges, by_type)
    return {"nodes": n_nodes, "edges": n_edges, "nodes_by_type": by_type}


def export(store: Store, out_dir: Path) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    np = out_dir / "kg_nodes.parquet"
    ep = out_dir / "kg_edges.parquet"
    store.con.execute(f"COPY kg_nodes TO '{np}' (FORMAT PARQUET)")
    store.con.execute(f"COPY kg_edges TO '{ep}' (FORMAT PARQUET)")

    # Neo4j-style Cypher dump (capped for portability)
    cypher = out_dir / "kg_import.cypher"
    rows_n = store.con.execute(
        "SELECT node_id, node_type, label FROM kg_nodes").fetchall()
    rows_e = store.con.execute(
        "SELECT src, rel, dst FROM kg_edges").fetchall()
    with cypher.open("w") as fh:
        fh.write("// SPID knowledge graph — Neo4j import\n")
        for nid, ntype, label in rows_n:
            safe = (label or "").replace("'", "\\'")[:120]
            fh.write(f"MERGE (n:{ntype} {{id:'{nid}'}}) SET n.label='{safe}';\n")
        for src, rel, dst in rows_e:
            fh.write(f"MATCH (a {{id:'{src}'}}),(b {{id:'{dst}'}}) "
                     f"MERGE (a)-[:{rel.upper()}]->(b);\n")
    log.info("Exported KG: %s, %s, %s", np.name, ep.name, cypher.name)
    return {"nodes_parquet": str(np), "edges_parquet": str(ep), "cypher": str(cypher)}