#!/usr/bin/env python3
"""
Novyx Core - Query Engine v1.1.0 (tools/query.py)
-------------------------------------------------
Responsibility:
1. Semantic Search: Find artifacts by meaning using vector similarity.
2. Authority Search: Filter artifacts by external links (GitHub, LinkedIn).
3. Knowledge Graph Traversal: Explore connections between nodes.
4. Schema v1.1.0 Compatible (Handles 'novyx:' namespace).
"""

import argparse
import json
import sys
from pathlib import Path
from typing import List, Dict

# ML Dependency Handling
try:
    from sentence_transformers import SentenceTransformer, util
    import torch
    ML_AVAILABLE = True
except ImportError:
    ML_AVAILABLE = False

MEMORY_DIR = Path("memory")
MODEL_NAME = 'all-MiniLM-L6-v2'
DEFAULT_BATCH_SIZE = 100  # For streaming operations


class QueryEngine:
    def __init__(self, strict: bool = False, force_no_embeddings: bool = False):
        self.model = None
        self.artifacts = []
        self.strict = strict
        self.force_no_embeddings = force_no_embeddings
        self.load_failures = 0
        self._load_memory()

        if self.force_no_embeddings:
            print("SEMANTIC MODE: OFF (--no-embeddings requested)")
        elif ML_AVAILABLE and self.artifacts:
            print("SEMANTIC MODE: ON (embeddings enabled)")
            # Lazy load model only if needed
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        elif not ML_AVAILABLE:
            print(
                "SEMANTIC MODE: OFF (missing 'sentence-transformers' or 'torch'). "
                "Install with: pip install sentence-transformers torch"
            )
        else:
            print("SEMANTIC MODE: OFF (no embeddings found)")

    def _load_memory(self):
        """Loads all JSON-LD artifacts into memory."""
        if not MEMORY_DIR.exists():
            return

        for path in MEMORY_DIR.rglob("*.jsonld"):
            try:
                with open(path, 'r') as f:
                    data = json.load(f)
                    # Normalize v1.0 and v1.1.0 fields
                    content_raw = data.get(
                        'articleBody') or data.get('content')

                    # Handle different content formats
                    if isinstance(content_raw, str):
                        content = content_raw
                    elif isinstance(content_raw, dict):
                        # Extract text from structured content (v1.0)
                        content = data.get('schema:description', '') or str(
                            content_raw.get('synthesis', ''))
                    else:
                        content = ''

                    embedding = data.get(
                        'novyx:embedding') or data.get('embedding')
                    links = data.get('novyx:semanticLinks') or data.get(
                        'semanticLinks', [])

                    self.artifacts.append({
                        'uuid': data.get('uuid') or data.get('Persistent_ID') or data.get('id'),
                        'name': data.get('name') or data.get('schema:name', 'Untitled'),
                        'content': content,
                        'embedding': embedding,
                        'links': links,
                        'sameAs': data.get('sameAs', []),
                        'path': str(path),
                        'category': path.parent.name
                    })
            except Exception as e:
                self.load_failures += 1
                message = f"⚠️  Warning: Failed to load artifact {path}: {type(e).__name__}: {e}"
                print(message, file=sys.stderr)
                if self.strict:
                    raise
                continue
        if self.load_failures:
            print(f"⚠️  Warning: Skipped {self.load_failures} artifact(s) due to load errors.", file=sys.stderr)
        # print(f"🧠 Knowledge Graph Loaded: {len(self.artifacts)} nodes")

    def search_semantic(self, query: str, top_k: int = 3, min_score: float = 0.2, json_output: bool = False):
        """Performs deep semantic search using vector embeddings.
        
        Args:
            query: Search query text
            top_k: Number of top results to return
            min_score: Minimum similarity score (0.0-1.0)
            json_output: If True, return results as JSON instead of printing
        
        Returns:
            List of result dicts if json_output=True, None otherwise
        """
        if self.force_no_embeddings:
            if json_output:
                return {"error": "Semantic search unavailable (--no-embeddings requested)"}
            print("❌ Semantic search unavailable (--no-embeddings requested).")
            return None
        if not ML_AVAILABLE:
            if json_output:
                return {"error": "Semantic search unavailable (missing dependencies)"}
            print("❌ Semantic search unavailable (missing dependencies).")
            return None

        if not self.model:
            if not json_output:
                print(f"Loading {MODEL_NAME}...")
            self.model = SentenceTransformer(MODEL_NAME, device=self.device)

        # Encode query
        query_embedding = self.model.encode(query, convert_to_tensor=True)

        # Prepare corpus with valid embeddings
        corpus_artifacts = [a for a in self.artifacts if a['embedding']]
        if not corpus_artifacts:
            if json_output:
                return {"error": "No vector embeddings found in memory", "results": []}
            print("⚠️  No vector embeddings found in memory.")
            return None

        corpus_embeddings = torch.tensor([a['embedding'] for a in corpus_artifacts]).to(self.device)

        # Search
        hits = util.semantic_search(
            query_embedding, corpus_embeddings, top_k=top_k)

        # Process results
        results = []
        for hit in hits[0]:
            if hit['score'] >= min_score:
                idx = hit['corpus_id']
                artifact = corpus_artifacts[idx]
                results.append({
                    'uuid': artifact['uuid'],
                    'name': artifact['name'],
                    'score': float(hit['score']),
                    'category': artifact['category'],
                    'excerpt': artifact['content'][:200] if artifact['content'] else '',
                    'path': artifact['path'],
                    'authority': artifact['sameAs'][0] if artifact['sameAs'] else None
                })

        if json_output:
            return {
                "query": query,
                "total_results": len(results),
                "min_score": min_score,
                "results": results
            }

        # Pretty print results
        print(f"\n🔍 Semantic Results for: '{query}'")
        print("-" * 50)

        if results:
            for result in results:
                print(f"📄 {result['name']} (Score: {result['score']:.4f})")
                print(f"   📂 Category: {result['category']}")
                print(f"   💡 Excerpt: {result['excerpt']}...")
                if result['authority']:
                    print(f"   🌐 Authority: {result['authority']}")
                print("")
        else:
            print("No relevant artifacts found (try lowering --min-score).")
        
        return results

    def search_semantic_streaming(
        self,
        query: str,
        directory: Path = None,
        *,
        top_k: int = 10,
        min_score: float = 0.2,
        batch_size: int = DEFAULT_BATCH_SIZE,
        tenant_id: str = None
    ) -> List[Dict]:
        """
        Streaming semantic search for large datasets (Phase 13).

        Processes artifacts in batches to limit memory usage.
        Maintains a running top-K across all batches.

        Args:
            query: Search query text
            directory: Directory to search (default: MEMORY_DIR)
            top_k: Number of top results to return
            min_score: Minimum similarity score
            batch_size: Artifacts per batch (default: 100)
            tenant_id: Optional tenant filter

        Returns:
            List of top-K results across entire dataset

        Example:
            >>> results = engine.search_semantic_streaming(
            ...     "deployment strategy",
            ...     top_k=5,
            ...     batch_size=50
            ... )
        """
        if self.force_no_embeddings or not ML_AVAILABLE:
            return []

        try:
            from tools.batch import stream_artifacts
        except ImportError:
            # Fall back to standard search
            return self.search_semantic(query, top_k=top_k, min_score=min_score, json_output=True).get('results', [])

        if not self.model:
            self.model = SentenceTransformer(MODEL_NAME, device=self.device)

        # Encode query once
        query_embedding = self.model.encode(query, convert_to_tensor=True)

        # Running top-K results
        all_results = []
        search_dir = directory or MEMORY_DIR

        for batch in stream_artifacts(search_dir, batch_size=batch_size, tenant_id=tenant_id):
            # Filter to artifacts with embeddings
            batch_with_embeddings = []
            for artifact in batch:
                embedding = artifact.get('novyx:embedding') or artifact.get('embedding')
                if embedding:
                    batch_with_embeddings.append({
                        'data': artifact,
                        'embedding': embedding
                    })

            if not batch_with_embeddings:
                continue

            # Compute similarities for this batch
            batch_embeddings = torch.tensor([a['embedding'] for a in batch_with_embeddings]).to(self.device)
            hits = util.semantic_search(query_embedding, batch_embeddings, top_k=min(top_k, len(batch_with_embeddings)))

            # Add qualifying results
            for hit in hits[0]:
                if hit['score'] >= min_score:
                    idx = hit['corpus_id']
                    artifact = batch_with_embeddings[idx]['data']
                    content = artifact.get('articleBody') or artifact.get('content', '')
                    if isinstance(content, dict):
                        content = str(content.get('synthesis', ''))

                    all_results.append({
                        'uuid': artifact.get('uuid') or artifact.get('Persistent_ID'),
                        'name': artifact.get('name') or artifact.get('schema:name', 'Untitled'),
                        'score': float(hit['score']),
                        'category': artifact.get('_path', '').split('/')[-2] if artifact.get('_path') else '',
                        'excerpt': content[:200] if content else '',
                        'path': artifact.get('_path', ''),
                        'tenant': artifact.get('_tenant')
                    })

        # Sort all results and return top-K
        all_results.sort(key=lambda x: x['score'], reverse=True)
        return all_results[:top_k]

    def search_authority(self, domain: str):
        """Finds artifacts linked to a specific external authority."""
        print(f"\n🌐 Searching for External Authority: '{domain}'")
        print("-" * 50)

        count = 0
        for art in self.artifacts:
            for link in art['sameAs']:
                if domain.lower() in link.lower():
                    print(f"✅ Found in: {art['name']}")
                    print(f"   🔗 Link: {link}")
                    count += 1

        if count == 0:
            print("No artifacts linked to this authority.")
        else:
            print(f"\nTotal matches: {count}")

    def show_stats(self):
        """Displays knowledge graph statistics."""
        print(f"\n📊 Knowledge Graph Statistics")
        print("-" * 30)
        print(f"Total Artifacts: {len(self.artifacts)}")

        # Count by category
        cats = {}
        links_count = 0
        authority_count = 0

        for a in self.artifacts:
            c = a['category']
            cats[c] = cats.get(c, 0) + 1
            links_count += len(a['links'])
            if a['sameAs']:
                authority_count += 1

        for c, count in cats.items():
            print(f"  • {c}: {count}")

        print("-" * 30)
        print(f"🧠 Semantic Connections: {links_count}")
        print(f"🌐 Authority Links:      {authority_count}")
        print(f"🛡️  Integrity Status:     Active")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Novyx Core Query Engine v1.2.0 (Phase 8: CLI Embeddings)")
    parser.add_argument("--search", help="Semantic search query")
    parser.add_argument("--semantic", action="store_true",
                        help="Enable semantic (embedding-based) search mode")
    parser.add_argument(
        "--authority", help="Filter by external authority (e.g., 'github', 'linkedin')")
    parser.add_argument("--stats", action="store_true",
                        help="Show knowledge graph stats")
    parser.add_argument("--min-score", type=float, default=0.2,
                        help="Minimum similarity score for results (0.0-1.0, default: 0.2)")
    parser.add_argument("--top-k", type=int, default=10,
                        help="Number of top results to return (default: 10)")
    parser.add_argument("--json", action="store_true", dest="json_output",
                        help="Output results as JSON")
    parser.add_argument("--strict", action="store_true",
                        help="Fail fast on artifact load errors")
    parser.add_argument("--no-embeddings", action="store_true",
                        help="Disable embeddings even if dependencies are installed")

    args = parser.parse_args()

    engine = QueryEngine(strict=args.strict, force_no_embeddings=args.no_embeddings)

    if args.search:
        results = engine.search_semantic(
            args.search, 
            top_k=args.top_k,
            min_score=args.min_score,
            json_output=args.json_output
        )
        if args.json_output and results:
            print(json.dumps(results, indent=2))
    elif args.authority:
        engine.search_authority(args.authority)
    elif args.stats:
        engine.show_stats()
    else:
        parser.print_help()
