"""
APOC trigger installation for tracking graph updates.

These triggers record changes to the graph in ChangeLog nodes that can be
polled by the updates server.
"""

import os
from typing import Optional

import dotenv
import neo4j

dotenv.load_dotenv()


def install_triggers(
    neo4j_uri: Optional[str] = None,
    neo4j_username: Optional[str] = None,
    neo4j_password: Optional[str] = None,
    database_name: str = "neo4j",
    updates_server_url: str = "https://localhost:10086",
    install_notify_trigger: bool = True,
) -> None:
    """
    Installs APOC triggers for tracking graph updates.

    Creates triggers that log:
    - Node creations (Message nodes)
    - Relationship creations (DEPENDS_ON edges)
    - Node deletions
    - Relationship deletions

    Each change is recorded as a ChangeLog node with a sequence number
    for ordered retrieval by the updates server.

    Optionally installs a notification trigger that calls the updates
    server's /notify endpoint when any relevant change occurs.

    Args:
        install_notify_trigger: If True, install a trigger that notifies
            the updates server via HTTP when changes occur. This requires
            Neo4j to have network access to the updates server URL.
            If False, the updates server will rely on polling.

    Currently, the observability graph does not have an external update
    API for existing nodes, so we won't install the associated triggers.
    """
    uri = neo4j_uri or os.environ["NEO4J_URI"]
    auth = (
        neo4j_username or os.environ["NEO4J_USERNAME"],
        neo4j_password or os.environ["NEO4J_PASSWORD"],
    )

    with neo4j.GraphDatabase.driver(uri, auth=auth) as driver:
        driver.verify_connectivity()

        # Create index for efficient ChangeLog queries
        try:
            driver.execute_query(
                """
                CREATE INDEX changelog_sequence IF NOT EXISTS
                FOR (c:ChangeLog)
                ON (c.sequence)
                """,
                database_=database_name,
            )
        except neo4j.exceptions.ClientError:
            pass

        # Create a sequence counter node if it doesn't exist
        driver.execute_query(
            """
            MERGE (s:ChangeLogSequence {id: 'main'})
            ON CREATE SET s.counter = 0
            """,
            database_=database_name,
        )

        # Install trigger for node creations
        # Use coalesce for optional properties to avoid missing property warnings
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_node_create',
                '
                UNWIND $createdNodes AS n
                WITH n WHERE n:Message
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH n, newValue AS seq
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "node_created",
                    id: n.id,
                    content: n.content,
                    role: n.role,
                    agent: n.agent,
                    tools: coalesce(n.tools, "[]"),
                    derived_from: n.derived_from,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for relationship creations
        # Use coalesce for optional edge properties
        # Use apoc.atomic.add for deadlock-free counter increment
        # Neo4j: (dst)-[:DEPENDS_ON]->(src) means dst depends on src
        # startNode = dst (dependent), endNode = src (dependency)
        # ChangeLog should have source_id = dependency, destination_id = dependent
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_rel_create',
                '
                UNWIND $createdRelationships AS r
                WITH r WHERE type(r) = "DEPENDS_ON"
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH r, newValue AS seq, startNode(r) AS dependent, endNode(r) AS dependency
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "edge_created",
                    source_id: dependency.id,
                    destination_id: dependent.id,
                    message_index: r.message_index,
                    proximal: r.proximal,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for node deletions
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_node_delete',
                '
                UNWIND $deletedNodes AS n
                WITH n WHERE n:Message
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH n, newValue AS seq
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "node_deleted",
                    id: n.id,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for relationship deletions
        # Use apoc.atomic.add for deadlock-free counter increment
        # Same semantics as edge creation: source_id = dependency, destination_id = dependent
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_rel_delete',
                '
                UNWIND $deletedRelationships AS r
                WITH r WHERE type(r) = "DEPENDS_ON"
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH r, newValue AS seq, startNode(r) AS dependent, endNode(r) AS dependency
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "edge_deleted",
                    source_id: dependency.id,
                    destination_id: dependent.id,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # ============================================
        # COMPUTATION (OTel Span) TRIGGERS
        # ============================================

        # Install trigger for Computation node creations
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_computation_create',
                '
                UNWIND $createdNodes AS n
                WITH n WHERE n:Computation
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH n, newValue AS seq
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "computation_created",
                    span_id: n.span_id,
                    trace_id: n.trace_id,
                    parent_span_id: n.parent_span_id,
                    name: n.name,
                    start_time_ns: n.start_time_ns,
                    end_time_ns: n.end_time_ns,
                    duration_ns: n.duration_ns,
                    status_code: n.status_code,
                    status_message: n.status_message,
                    attributes: coalesce(n.attributes, "{}"),
                    events: coalesce(n.events, "[]"),
                    service_name: n.service_name,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for Computation node deletions
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_computation_delete',
                '
                UNWIND $deletedNodes AS n
                WITH n WHERE n:Computation
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH n, newValue AS seq
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "computation_deleted",
                    span_id: n.span_id,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for CHILD_OF relationship creations
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_child_of_create',
                '
                UNWIND $createdRelationships AS r
                WITH r WHERE type(r) = "CHILD_OF"
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH r, newValue AS seq, startNode(r) AS child, endNode(r) AS parent
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "computation_edge_created",
                    child_span_id: child.span_id,
                    parent_span_id: parent.span_id,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for PRODUCES relationship creations
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_produces_create',
                '
                UNWIND $createdRelationships AS r
                WITH r WHERE type(r) = "PRODUCES"
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH r, newValue AS seq, startNode(r) AS comp, endNode(r) AS msg
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "computation_message_link_created",
                    span_id: comp.span_id,
                    message_id: msg.id,
                    edge_type: "PRODUCES",
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Install trigger for CONSUMES relationship creations
        # Use apoc.atomic.add for deadlock-free counter increment
        driver.execute_query(
            """
            CALL apoc.trigger.install(
                $database,
                'log_consumes_create',
                '
                UNWIND $createdRelationships AS r
                WITH r WHERE type(r) = "CONSUMES"
                MATCH (s:ChangeLogSequence {id: "main"})
                CALL apoc.atomic.add(s, "counter", 1) YIELD newValue
                WITH r, newValue AS seq, startNode(r) AS comp, endNode(r) AS msg
                CREATE (c:ChangeLog {
                    sequence: seq,
                    type: "computation_message_link_created",
                    span_id: comp.span_id,
                    message_id: msg.id,
                    edge_type: "CONSUMES",
                    message_index: r.message_index,
                    timestamp: datetime()
                })
                ',
                {phase: 'afterAsync'}
            )
            """,
            database_="system",
            database=database_name,
        )

        # Optionally install trigger to notify updates server on any relevant change
        # This trigger fires on all events and makes an HTTP POST to /notify
        # Note: URL is embedded directly because APOC params aren't Cypher params
        if install_notify_trigger:
            notify_url = f"{updates_server_url}/notify"
            # Updated to include Computation nodes and their relationships
            trigger_query = f"""
                CALL apoc.trigger.install(
                    $database,
                    'notify_updates_server',
                    '
                    WITH $createdNodes AS cn, $deletedNodes AS dn,
                         $createdRelationships AS cr, $deletedRelationships AS dr
                    WITH [n IN cn WHERE n:Message] AS msg_created,
                         [n IN dn WHERE n:Message] AS msg_deleted,
                         [r IN cr WHERE type(r) = "DEPENDS_ON"] AS rel_created,
                         [r IN dr WHERE type(r) = "DEPENDS_ON"] AS rel_deleted,
                         [n IN cn WHERE n:Computation] AS comp_created,
                         [n IN dn WHERE n:Computation] AS comp_deleted,
                         [r IN cr WHERE type(r) IN ["CHILD_OF", "PRODUCES", "CONSUMES"]] AS comp_rel_created
                    WITH msg_created + msg_deleted + rel_created + rel_deleted +
                         comp_created + comp_deleted + comp_rel_created AS changes
                    WHERE size(changes) > 0
                    CALL apoc.load.jsonParams("{notify_url}", {{}}, null, "", {{method: "POST"}})
                    YIELD value
                    RETURN value
                    ',
                    {{phase: "afterAsync"}}
                )
                """
            driver.execute_query(
                trigger_query,
                database_="system",
                database=database_name,
            )
            print(f"  Notification URL: {notify_url}")
        else:
            # Uninstall existing notify trigger if present
            # First check what triggers exist
            try:
                show_result = driver.execute_query(
                    "CALL apoc.trigger.show($database) YIELD name RETURN name",
                    database_="system",
                    database=database_name,
                )
                existing_triggers = [r["name"] for r in show_result.records]
                print(f"  Existing triggers: {existing_triggers}")

                if "notify_updates_server" in existing_triggers:
                    driver.execute_query(
                        "CALL apoc.trigger.drop($database, $name)",
                        database_="system",
                        database=database_name,
                        name="notify_updates_server",
                    )
                    print("  Removed notification trigger (polling mode)")
                else:
                    print("  Notification trigger not installed (polling mode)")
            except neo4j.exceptions.ClientError as e:
                print(f"  Warning: Could not manage triggers: {e}")

        print(f"APOC triggers installed successfully for database '{database_name}'")


def uninstall_triggers(
    neo4j_uri: Optional[str] = None,
    neo4j_username: Optional[str] = None,
    neo4j_password: Optional[str] = None,
    database_name: str = "neo4j",
) -> None:
    """Removes all installed APOC triggers for graph updates."""
    uri = neo4j_uri or os.environ["NEO4J_URI"]
    auth = (
        neo4j_username or os.environ["NEO4J_USERNAME"],
        neo4j_password or os.environ["NEO4J_PASSWORD"],
    )

    with neo4j.GraphDatabase.driver(uri, auth=auth) as driver:
        driver.verify_connectivity()

        for trigger_name in [
            "log_node_create",
            "log_rel_create",
            "log_node_delete",
            "log_rel_delete",
            "log_computation_create",
            "log_computation_delete",
            "log_child_of_create",
            "log_produces_create",
            "log_consumes_create",
            "notify_updates_server",
        ]:
            try:
                driver.execute_query(
                    "CALL apoc.trigger.drop($database, $name)",
                    database_="system",
                    database=database_name,
                    name=trigger_name,
                )
            except neo4j.exceptions.ClientError:
                pass

        print(f"APOC triggers removed for database '{database_name}'")


if __name__ == "__main__":
    from fire import Fire  # type: ignore[import-untyped]

    Fire({"install": install_triggers, "uninstall": uninstall_triggers})
