#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Graph RAG Concept Extractor"""
import json
from pulse.workers.worker_llm import WorkerLLM

SCANNER_PROVIDER = "ollama"
SCANNER_MODEL = "qwen3-coder-next:cloud"


def extract_graph_rag_concepts(query: str) -> tuple[str, str] | None:
    """Use a fast LLM to find the two core concepts for a causal chain query."""
    if not query:
        return None

    llm = WorkerLLM(SCANNER_PROVIDER, SCANNER_MODEL)
    llm.init_client()

    sys_ctx = (
        "You are a linguistic concept extractor. Your job is to read a user's task "
        "and identify the two most important nouns or concepts to use for a causal graph search. "
        "Output ONLY a JSON object with two keys: 'concept1' and 'concept2'.\n\n"
        "EXAMPLES:\n"
        "User: 'Fix the Docker bug that causes memory leaks'\n"
        'JSON: {"concept1": "Docker", "concept2": "memory leak"}\n'
        "User: 'Refactor the API to improve security'\n"
        'JSON: {"concept1": "API", "concept2": "security"}\n'
    )
    user_msg = f"User: '{query}'\nJSON:"

    try:
        text, _, _ = llm.call_llm(sys_ctx, [{"role": "user", "content": user_msg}])
        text = text.replace("```json", "").replace("```", "").strip()
        concepts = json.loads(text)
        c1 = concepts.get("concept1")
        c2 = concepts.get("concept2")
        if c1 and c2:
            return (c1, c2)
    except Exception:
        pass

    # Fallback to simple keyword split
    parts = query.split()
    return (parts[0], parts[-1]) if len(parts) >= 2 else None
