memory/graph_memory.py

"""
Graph-RAG Memory: stores and retrieves knowledge from Neo4j.
This is the core solution to the LLM context window problem.
Each agent writes results as graph nodes; retrieval fetches only
the relevant subgraph, keeping LLM prompts small (~3-5k tokens).

Enhanced with hybrid search: combines Cypher keyword matching (precision)
with vector similarity search (recall) for better claim discovery.
"""

from typing import Any, Dict, List, Optional
import json as _json

from memory.base_memory import BaseMemory
from services.neo4j_service import Neo4jService
from services.vector_store_service import VectorStoreService
from services.lancedb_service import LanceDBService
from utils.logger import get_logger
from utils.text import format_context_block, count_tokens

logger = get_logger("graph_memory")

MAX_CONTEXT_TOKENS = 5000


class GraphMemory(BaseMemory):
    """
    Neo4j-backed Graph-RAG memory with hybrid search.
    Provides focused context retrieval using both keyword matching and semantic similarity.
    Reduces LLM context window usage while improving claim discovery.
    """

    def __init__(
        self,
        neo4j: Neo4jService,
        embedder: VectorStoreService,
        vector_db: Optional[LanceDBService] = None,
    ):
        self.neo4j = neo4j
        self.embedder = embedder
        self.vector_db = vector_db

    # ------------------------------------------------------------------ #
    # BaseMemory interface
    # ------------------------------------------------------------------ #

    async def store(
        self, key: str, value: Any, metadata: Optional[Dict] = None
    ) -> None:
        """Generic store: saves a key-value pair as a (:Memory) node."""
        meta = metadata or {}
        self.neo4j.run_write(
            """
            MERGE (m:Memory {key: $key})
            SET m.value = $value, m.job_id = $job_id
            """,
            key=key,
            value=str(value),
            job_id=meta.get("job_id", ""),
        )

    async def retrieve(self, key: str) -> Optional[Any]:
        rows = self.neo4j.run(
            "MATCH (m:Memory {key: $key}) RETURN m.value AS v", key=key
        )
        return rows[0]["v"] if rows else None

    async def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Keyword search across Memory nodes (fallback without vector index)."""
        rows = self.neo4j.run(
            """
            MATCH (m:Memory)
            WHERE toLower(m.value) CONTAINS toLower($q)
            RETURN m.key AS key, m.value AS value
            LIMIT $k
            """,
            q=query,
            k=k,
        )
        return rows

    async def delete(self, key: str) -> None:
        self.neo4j.run_write("MATCH (m:Memory {key: $key}) DETACH DELETE m", key=key)

    async def clear(self, scope: Optional[str] = None) -> None:
        if scope:
            self.neo4j.run_write(
                "MATCH (m:Memory {job_id: $scope}) DETACH DELETE m", scope=scope
            )
        else:
            self.neo4j.run_write("MATCH (m:Memory) DETACH DELETE m")

    # ------------------------------------------------------------------ #
    # Graph-RAG: domain-specific context assembly
    # ------------------------------------------------------------------ #

    def get_sections_context(
        self, doc_id: str, headings: Optional[List[str]] = None
    ) -> str:
        """Retrieve section texts as a focused LLM context block."""
        if headings:
            rows = self.neo4j.run(
                """
                MATCH (d:Document {doc_id: $doc_id})-[:HAS_SECTION]->(s:Section)
                WHERE toLower(s.heading) IN $headings
                RETURN s.heading AS heading, s.text AS text
                ORDER BY s.position
                """,
                doc_id=doc_id,
                headings=[h.lower() for h in headings],
            )
        else:
            rows = self.neo4j.run(
                """
                MATCH (d:Document {doc_id: $doc_id})-[:HAS_SECTION]->(s:Section)
                RETURN s.heading AS heading, s.text AS text
                ORDER BY s.position
                """,
                doc_id=doc_id,
            )

        blocks, tokens = [], 0
        for row in rows:
            block = format_context_block(row["heading"], row["text"] or "")
            cost = count_tokens(block)
            if tokens + cost > MAX_CONTEXT_TOKENS:
                break
            blocks.append(block)
            tokens += cost

        logger.debug(f"Sections context: {len(blocks)} blocks, ~{tokens} tokens")
        return "\n".join(blocks)

    def hybrid_search_sections(
        self,
        doc_id: str,
        query: str,
        top_k: int = 10,
        keyword_weight: float = 0.6,
        vector_weight: float = 0.4,
    ) -> str:
        """
        Hybrid search combining keyword matching (Cypher) and vector similarity (LanceDB).

        Strategy:
        1. Extract keywords from query and search Neo4j sections (high precision)
        2. Encode query and search vector DB (high recall, semantic)
        3. Combine results, deduplicate, and rank by weighted score
        4. Format as context blocks within token budget

        Args:
            doc_id: Document ID to search within
            query: Natural language query
            top_k: Number of sections to return
            keyword_weight: Weight for keyword match scores (0-1)
            vector_weight: Weight for vector similarity scores (0-1)

        Returns:
            Formatted context string with top sections
        """

        results_dict = {}  # section_id -> {heading, text, score, method}

        # METHOD 1: Keyword search in Neo4j (precision)
        try:
            # Extract keywords from query (simple: split and filter common words)
            keywords = [
                w
                for w in query.lower().split()
                if len(w) > 3
                and w not in {"the", "this", "that", "with", "from", "into"}
            ][:5]

            if keywords:
                # Cypher query using text containment
                cypher_query = """
                MATCH (d:Document {doc_id: $doc_id})-[:HAS_SECTION]->(s:Section)
                WHERE toLower(s.text) CONTAINS any(kw in $keywords WHERE toLower(s.text) CONTAINS kw)
                   OR toLower(s.heading) CONTAINS any(kw in $keywords WHERE toLower(s.heading) CONTAINS kw)
                RETURN s.section_id AS section_id,
                       s.heading AS heading,
                       s.text AS text,
                       s.position AS position
                """

                rows = self.neo4j.run(
                    cypher_query,
                    doc_id=doc_id,
                    keywords=keywords,
                )

                for row in rows:
                    section_id = row.get("section_id")
                    if section_id:
                        results_dict[section_id] = {
                            "heading": row.get("heading", "Unknown"),
                            "text": row.get("text", ""),
                            "position": row.get("position", 999),
                            "score": keyword_weight,
                            "method": "keyword",
                        }

                logger.debug(
                    f"Keyword search found {len([r for r in results_dict.values() if r['method'] == 'keyword'])} sections"
                )
        except Exception as e:
            logger.warning(f"Keyword search failed: {e}")

        # METHOD 2: Vector similarity search in LanceDB (recall)
        if self.vector_db:
            table_name = f"doc_{doc_id}_sections"
            if self.vector_db.table_exists(table_name):
                try:
                    query_vector = self.embedder.embed(query)
                    vector_results = self.vector_db.search(
                        table_name, query_vector, top_k=top_k * 2
                    )
                    for result in vector_results:
                        section_id = result.get("id")
                        if section_id:
                            if section_id in results_dict:
                                current_score = results_dict[section_id]["score"]
                                vector_score = float(result.get("_distance", 0.0))
                                combined_score = current_score * 0.7 + vector_score * 0.3
                                results_dict[section_id]["score"] = combined_score
                                results_dict[section_id]["method"] = "hybrid"
                            else:
                                vector_score = float(result.get("_distance", 0.0))
                                results_dict[section_id] = {
                                    "heading": result.get("heading", "Unknown"),
                                    "text": result.get("text", ""),
                                    "position": 999,
                                    "score": vector_weight,
                                    "method": "vector",
                                }
                    logger.debug(
                        f"Vector search found {len([r for r in results_dict.values() if r['method'] in ('vector', 'hybrid')])} sections"
                    )
                except Exception as e:
                    logger.warning(f"Vector search failed: {e}")
            else:
                logger.debug(f"Vector table {table_name} not found — skipping vector search")

        # STEP 3: Rank and format results
        if not results_dict:
            logger.debug("Hybrid search returned no results")
            return ""

        # Sort by score descending, then by position ascending
        sorted_results = sorted(
            results_dict.values(), key=lambda x: (-x["score"], x["position"])
        )

        # Build context blocks within token budget
        blocks = []
        tokens = 0

        for result in sorted_results[:top_k]:
            block = format_context_block(result["heading"], result["text"] or "")
            cost = count_tokens(block)

            if tokens + cost > MAX_CONTEXT_TOKENS:
                break

            blocks.append(block)
            tokens += cost

        logger.debug(
            f"Hybrid search: {len(blocks)} blocks, ~{tokens} tokens, "
            f"methods={set(r['method'] for r in sorted_results[:top_k])}"
        )

        return "\n".join(blocks)

    def get_claims_context(self, doc_id: str, limit: int = 15) -> str:
        """Retrieve top claims as context for argumentation/scoring agents."""
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_SECTION]->(:Section)
                  -[:CONTAINS_CLAIM]->(c:Claim)
            RETURN c.text AS text, c.type AS type,
                   c.confidence AS confidence, c.composite_score AS score
            ORDER BY c.composite_score DESC
            LIMIT $limit
            """,
            doc_id=doc_id,
            limit=limit,
        )
        if not rows:
            return ""
        lines = [
            f"- [{r['type']}] {r['text']} (confidence={r['confidence']:.2f})"
            for r in rows
        ]
        return format_context_block("Extracted Claims", "\n".join(lines))

    def get_keywords_context(self, doc_id: str) -> str:
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_SECTION]->(s:Section)
                  -[:HAS_KEYWORD]->(k:Keyword)
            RETURN DISTINCT k.term AS term, k.score AS score
            ORDER BY k.score DESC
            LIMIT 30
            """,
            doc_id=doc_id,
        )
        terms = [r["term"] for r in rows]
        return format_context_block("Keywords", ", ".join(terms))

    def get_gaps_context(self, doc_id: str) -> str:
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_GAP]->(g:CitationGap)
            OPTIONAL MATCH (g)-[:SUGGESTS]->(p:SuggestedPaper)
            RETURN g.description AS desc, g.severity AS severity,
                   collect(p.title) AS suggestions
            ORDER BY g.severity DESC
            LIMIT 20
            """,
            doc_id=doc_id,
        )
        if not rows:
            return ""
        lines = [
            f"- [{r['severity']}] {r['desc']} | Suggested: {', '.join(r['suggestions'][:2])}"
            for r in rows
        ]
        return format_context_block("Citation Gaps", "\n".join(lines))

    def get_counterfactuality_context(self, doc_id: str) -> str:
        """Retrieve counterfactuality analysis as context."""
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_COUNTERFACTUALITY]->(cf:CounterFactuality)
            RETURN cf.overall_score AS overall_score, cf.most_common_type AS common_type,
                   cf.concerns AS concerns, cf.strengths AS strengths, cf.summary AS summary
            """,
            doc_id=doc_id,
        )
        if not rows:
            return ""
        r = rows[0]
        concerns = r.get("concerns", []) or []
        strengths = r.get("strengths", []) or []
        lines = [
            f"Overall counterfactuality score: {r.get('overall_score', 0):.2f}/1.0",
            f"Most common type: {r.get('common_type', 'factual')}",
            f"Summary: {r.get('summary', '')}",
        ]
        if concerns:
            lines.append(f"Concerns: {'; '.join(concerns[:3])}")
        if strengths:
            lines.append(f"Strengths: {'; '.join(strengths[:3])}")

        claim_rows = self.neo4j.run(
            """
            MATCH (c:Claim {doc_id: $doc_id})
            WHERE c.counterfactual_type IS NOT NULL
            RETURN c.text AS text, c.counterfactual_type AS cf_type,
                   c.counterfactual_score AS cf_score
            ORDER BY c.counterfactual_score DESC
            LIMIT 10
            """,
            doc_id=doc_id,
        )
        if claim_rows:
            lines.append("\nPer-claim analysis (most counterfactual first):")
            for cr in claim_rows:
                lines.append(
                    f"  - [{cr.get('cf_type', 'unknown')}] "
                    f"(score={cr.get('cf_score', 0):.2f}) {cr.get('text', '')}"
                )

        return format_context_block("Counterfactuality Analysis", "\n".join(lines))

    def get_novelty_context(self, doc_id: str) -> str:
        """Retrieve problem-statement novelty analysis as context."""
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_NOVELTY]->(n:Novelty)
            RETURN n.novelty_score AS novelty_score, n.verdict AS verdict,
                   n.problem_statement AS problem_statement, n.rationale AS rationale,
                   n.similar_works AS similar_works
            """,
            doc_id=doc_id,
        )
        if not rows:
            return ""
        r = rows[0]
        lines = [
            f"Novelty score: {r.get('novelty_score', 0):.1f}/10 ({r.get('verdict', 'UNKNOWN')})",
            f"Problem statement: {r.get('problem_statement', '')}",
            f"Rationale: {r.get('rationale', '')}",
        ]
        works = r.get("similar_works", []) or []
        if isinstance(works, str):
            try:
                works = _json.loads(works) or []
            except Exception:
                works = []
        if works:
            lines.append("Closely related existing work:")
            for w in works[:5]:
                title = w.get("title", "Related work") if isinstance(w, dict) else str(w)
                year = w.get("year", "") if isinstance(w, dict) else ""
                lines.append(f"  - {title} ({year})")
        return format_context_block("Problem Novelty", "\n".join(lines))

    def get_citation_relevance_context(self, doc_id: str) -> str:
        """Retrieve off-topic citation flags as context."""
        rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_IRRELEVANT_CITATION]->(ic:IrrelevantCitation)
            RETURN ic.index AS index, ic.verdict AS verdict, ic.reason AS reason
            ORDER BY ic.index
            """,
            doc_id=doc_id,
        )
        if not rows:
            return ""
        lines = [
            f"- [{r['verdict']}] Ref [{r['index']}]: {r['reason']}" for r in rows[:10]
        ]
        return format_context_block("Citation Relevance", "\n".join(lines))

    def get_full_report_context(self, doc_id: str) -> str:
        """Assemble complete report context from graph (~6k tokens max)."""
        parts = [
            self.get_keywords_context(doc_id),
            self.get_claims_context(doc_id),
            self.get_counterfactuality_context(doc_id),
            self.get_gaps_context(doc_id),
            self.get_novelty_context(doc_id),
            self.get_citation_relevance_context(doc_id),
        ]
        # Add scores
        score_rows = self.neo4j.run(
            """
            MATCH (d:Document {doc_id: $doc_id})-[:HAS_SCORE]->(sc:Score)
            RETURN sc.literary_score AS literary, sc.argument_score AS argument,
                   sc.citation_completeness AS citations, sc.overall_score AS overall
            """,
            doc_id=doc_id,
        )
        if score_rows:
            r = score_rows[0]
            score_text = (
                f"Literary: {r['literary']:.2f}/10 | "
                f"Argument: {r['argument']:.2f}/10 | "
                f"Citations: {r['citations']:.2f}/10 | "
                f"Overall: {r['overall']:.2f}/10"
            )
            parts.append(format_context_block("Scores", score_text))

        context = "\n".join(p for p in parts if p)
        tokens = count_tokens(context)
        logger.debug(f"Full report context: ~{tokens} tokens")
        return context