from neo4j import GraphDatabase
from dotenv import load_dotenv
import os
import spacy
import sqlite3
import shutil
from pathlib import Path

original_db = Path("data/nasa_data.db")
runtime_db = Path("/tmp/nasa_data.db")

if not runtime_db.exists():
    shutil.copy(original_db, runtime_db)
    print(f"copied {original_db} → {runtime_db}")
else:
    print(f"usiing existing runtime DB at {runtime_db}")

nlp = spacy.load("en_core_web_sm")

load_dotenv()
NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USER = os.getenv("NEO4J_USERNAME")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")

print("connecting to Neo4j at:", NEO4J_URI)

driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
driver.verify_connectivity()

with driver.session() as session:
    result = session.run("RETURN 'Neo4j connection successful!' AS msg")
    print(result.single()["msg"])

def extract_entities(text):
    """return (entity_text, entity_type) from text using spacy NER"""
    doc = nlp(text)
    return [(ent.text, ent.label_) for ent in doc.ents]


def extract_keywords(text):
    """return important nouns/proper nouns from text"""
    doc = nlp(text)
    return [token.lemma_.capitalize() for token in doc if token.pos_ in ("NOUN", "PROPN") and len(token.text) > 3]


def load_data():
    """laoding projects, publications, links, PIs, authors and abstracts from SQLite"""
    db_path = "/tmp/nasa_data.db"
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()

    cur.execute("SELECT Task_id, [Project Title] FROM projects")
    proj = cur.fetchall()

    cur.execute("SELECT pmc_id, title FROM publications")
    pub = cur.fetchall()

    cur.execute("SELECT project_id, pmc_id, link_type, confidence FROM links")
    links = cur.fetchall()

    cur.execute("SELECT Task_id, PI_Name FROM projects")
    pis = cur.fetchall()

    cur.execute("SELECT pmc_id, authors FROM publications")
    authors = cur.fetchall()

    cur.execute("SELECT pmc_id, abstract FROM publications")
    abstracts = cur.fetchall()

    conn.close()
    return proj, pub, links, pis, authors, abstracts


def clear_old_data(session):
    """optionallly clear old nodes and relationships for a clean rebuild"""
    session.run("MATCH ()-[r:LINKS_TO]->() DELETE r")
    session.run("MATCH (t:Topic) DETACH DELETE t")
    session.run("MATCH (e:Entity) DETACH DELETE e")


def create_links(session, proj, pub, links):
    """create proj and pub nodes and LINKS_TO relationships"""
    # proj
    for proj_id, title in proj:
        session.run("""
            MERGE (p:Project {id: $id})
            SET p.title = $title
        """, id=str(proj_id), title=title)

    # pub
    for pmc_id, title in pub:
        session.run("""
            MERGE (pub:Publication {pmc_id: $pmc_id})
            SET pub.title = $title
        """, pmc_id=str(pmc_id).upper(), title=title)

    # LINKS_TO relationships
    for proj_id, pmc_id, link_type, conf in links:
        if conf < 0.5:
            continue
        session.run("""
            MATCH (p:Project {id: $proj_id})
            MATCH (pub:Publication {pmc_id: $pmc_id})
            MERGE (p)-[r:LINKS_TO]->(pub)
            SET r.type = $link_type,
                r.confidence = $confidence
        """, proj_id=str(proj_id), pmc_id=str(pmc_id).upper(), link_type=link_type, confidence=conf)


def add_people_relationships(session, pis, authors):
    """add person nodes and link them as HAS_PI or HAS_AUTHOR"""
    # PIs
    for proj_id, pi_name in pis:
        if not pi_name:
            continue
        session.run("MERGE (person:Person {name: $pi_name})", pi_name=pi_name)
        session.run("""
            MATCH (p:Project {id: $proj_id})
            MATCH (person:Person {name: $pi_name})
            MERGE (p)-[:HAS_PI]->(person)
        """, proj_id=str(proj_id), pi_name=pi_name)

    # authors
    for pmc_id, author_str in authors:
        if not author_str:
            continue
        for author in [a.strip() for a in author_str.split(",") if a.strip()]:
            session.run("MERGE (person:Person {name: $author})", author=author)
            session.run("""
                MATCH (pub:Publication {pmc_id: $pmc_id})
                MATCH (person:Person {name: $author})
                MERGE (pub)-[:HAS_AUTHOR]->(person)
            """, pmc_id=str(pmc_id).upper(), author=author)


def add_coauthor_rel(session):
    session.run("""
        MATCH (a:Person)<-[:HAS_AUTHOR]-(pub:Publication)-[:HAS_AUTHOR]->(b:Person)
        WHERE a <> b
        MERGE (a)-[r:CO_AUTHOR_WITH]-(b)
        ON CREATE SET r.count = 1
        ON MATCH SET r.count = r.count + 1
    """)


def add_pi_author_overlap(session):
    session.run("""
        MATCH (p:Project)-[:HAS_PI]->(pi:Person)
        MATCH (pub:Publication)-[:HAS_AUTHOR]->(pi)
        MERGE (p)-[:PI_IS_AUTHOR_OF]->(pub)
    """)


def add_project_similarity(session):
    session.run("""
        MATCH (p1:Project)-[:LINKS_TO]->(pub:Publication)<-[:LINKS_TO]-(p2:Project)
        WHERE p1 <> p2
        MERGE (p1)-[r:RELATED_TO_PROJECT]->(p2)
        ON CREATE SET r.shared_pubs = 1
        ON MATCH SET r.shared_pubs = r.shared_pubs + 1
    """)


def add_author_project_links(session):
    session.run("""
        MATCH (p:Project)-[:LINKS_TO]->(pub:Publication)-[:HAS_AUTHOR]->(a:Person)
        MERGE (a)-[r:WORKS_ON]->(p)
        ON CREATE SET r.via_publications = 1
        ON MATCH SET r.via_publications = r.via_publications + 1
    """)


def add_topic_nodes(session, projects, publications):
    for proj_id, title in projects:
        if not title:
            continue
        for kw in extract_keywords(title):
            session.run("""
                MERGE (t:Topic {name: $topic})
                WITH t
                MATCH (p:Project {id: $proj_id})
                MERGE (p)-[:HAS_TOPIC]->(t)
            """, topic=kw, proj_id=str(proj_id))

    for pmc_id, title in publications:
        if not title:
            continue
        for kw in extract_keywords(title):
            session.run("""
                MERGE (t:Topic {name: $topic})
                WITH t
                MATCH (pub:Publication {pmc_id: $pmc_id})
                MERGE (pub)-[:HAS_TOPIC]->(t)
            """, topic=kw, pmc_id=str(pmc_id).upper())


def add_publication_text(session, abstracts):
    for pmc_id, abstract in abstracts:
        if not abstract:
            continue
        session.run("""
            MATCH (pub:Publication {pmc_id: $pmc_id})
            SET pub.abstract = $abstract
        """, pmc_id=str(pmc_id).upper(), abstract=abstract)


def add_entities_from_text(session, abstracts):
    for pmc_id, abstract in abstracts:
        if not abstract:
            continue
        entities = extract_entities(abstract)
        for entity_text, entity_type in entities:
            session.run("""
                MERGE (e:Entity {name: $name, type: $type})
                WITH e
                MATCH (pub:Publication {pmc_id: $pmc_id})
                MERGE (pub)-[:MENTIONS_ENTITY]->(e)
            """, name=entity_text, type=entity_type, pmc_id=str(pmc_id).upper())

proj, pub, links, pis, authors, abstracts = load_data()

with driver.session() as session:
    clear_old_data(session)
    create_links(session, proj, pub, links)
    add_people_relationships(session, pis, authors)
    add_coauthor_rel(session)
    add_pi_author_overlap(session)
    add_project_similarity(session)
    add_author_project_links(session)
    add_topic_nodes(session, proj, pub)
    add_publication_text(session, abstracts)
    add_entities_from_text(session, abstracts)

driver.close()
print("data pushed to Neo4j successfully!")