ke.py

#!/usr/bin/env python3
"""Knowledge Engine CLI — self-refining knowledge base with entity extraction,
graph relationships, composite scoring, and query expansion."""

import argparse
import json
import sys
import os

sys.path.insert(0, os.path.dirname(__file__))

import db
import scraper
import properties_scraper
import entities
import smart_search


def cmd_init(args):
    db.init_db()
    print("Knowledge Engine initialized.")


def cmd_add_source(args):
    db.init_db()
    config = {}
    if args.selector:
        config["selector"] = args.selector

    db.add_source(
        type_=args.type,
        name=args.name,
        url=args.url,
        config=config
    )
    print(f"Added source: {args.name} ({args.type})")


def cmd_scrape(args):
    db.init_db()
    print("Scraping all sources...")
    total = scraper.scrape_all()
    print(f"Done. Added {total} new documents total.")

    if args.extract_entities:
        print("\nExtracting entities from new documents...")
        entities.process_all_documents()


def cmd_query(args):
    db.init_db()
    query = " ".join(args.query)

    if args.smart:
        results = smart_search.composite_search(
            query,
            limit=args.limit,
            min_relevance=args.min_relevance or 0.0,
            topic=args.topic,
            recency_boost=not args.no_recency,
            use_expansion=not args.no_expand,
        )

        if not results:
            print("No results found.")
            return

        print(f"\n{'='*60}")
        print(f"  Smart Search: {query}")
        print(f"  Found: {len(results)} documents")
        print(f"{'='*60}\n")

        for item in results:
            r = item["doc"]
            rel_bar = "█" * int(r["relevance_score"] * 10) + "░" * (10 - int(r["relevance_score"] * 10))
            print(f"  [{r['id']}] {r['title']}")
            print(f"       Composite: {item['composite']:.3f}  BM25: {item['bm25']:.3f}  Entities: +{item['entity_bonus']:.3f}  Recency: +{item['recency']:.3f}")
            print(f"       Relevance: [{rel_bar}] {r['relevance_score']:.2f}")
            try:
                if r["topic"]:
                    print(f"       Topic: {r['topic']}")
            except (IndexError, KeyError):
                pass
            keywords = json.loads(r["keywords"]) if r["keywords"] else []
            if keywords:
                print(f"       Keywords: {', '.join(keywords[:8])}")
            if r["url"]:
                print(f"       URL: {r['url']}")
            snippet = r["content"][:200].replace("\n", " ")
            print(f"       {snippet}...")
            print()
    else:
        min_rel = args.min_relevance or 0.0
        results = db.search(query, limit=args.limit, min_relevance=min_rel)

        if not results:
            print("No results found.")
            return

        print(f"\n{'='*60}")
        print(f"  Results for: {query}")
        print(f"  Found: {len(results)} documents")
        print(f"{'='*60}\n")

        for r in results:
            rel_bar = "█" * int(r["relevance_score"] * 10) + "░" * (10 - int(r["relevance_score"] * 10))
            print(f"  [{r['id']}] {r['title']}")
            print(f"       Relevance: [{rel_bar}] {r['relevance_score']:.2f}")
            keywords = json.loads(r["keywords"]) if r["keywords"] else []
            if keywords:
                print(f"       Keywords: {', '.join(keywords[:8])}")
            if r["url"]:
                print(f"       URL: {r['url']}")
            snippet = r["content"][:200].replace("\n", " ")
            print(f"       {snippet}...")
            print()


def cmd_feedback(args):
    db.init_db()
    relevant = args.relevant.lower() in ("true", "1", "yes", "y")
    db.record_feedback(args.id, args.context or "", relevant)
    status = "relevant" if relevant else "not relevant"
    print(f"Recorded feedback: document {args.id} marked as {status}")


def cmd_sources(args):
    db.init_db()
    sources = db.get_sources(enabled_only=False)
    if not sources:
        print("No sources configured.")
        return

    print(f"\n{'='*60}")
    print(f"  Knowledge Sources")
    print(f"{'='*60}\n")

    for s in sources:
        status = "✓" if s["enabled"] else "✗"
        last = s["last_scraped"] or "never"
        topic = s["topic"] if "topic" in s.keys() else "—"
        print(f"  [{status}] {s['name']}")
        print(f"       Type: {s['type']} | Topic: {topic} | URL: {s['url']}")
        print(f"       Last scraped: {last}")
        print()


def cmd_stats(args):
    db.init_db()
    stats = db.get_stats()
    conn = db.get_conn()

    # Extended stats
    entity_count = conn.execute("SELECT COUNT(*) FROM entities").fetchone()[0]
    link_count = conn.execute("SELECT COUNT(*) FROM document_links").fetchone()[0]
    synonym_count = conn.execute("SELECT COUNT(*) FROM synonyms").fetchone()[0]
    collection_count = conn.execute("SELECT COUNT(*) FROM collections").fetchone()[0]

    # Top entity types
    entity_types = conn.execute("""
        SELECT type, COUNT(*) as cnt FROM entities GROUP BY type ORDER BY cnt DESC
    """).fetchall()

    # Most connected docs
    connected = conn.execute("""
        SELECT d.id, d.title,
               (SELECT COUNT(*) FROM document_links WHERE source_doc_id = d.id OR target_doc_id = d.id) AS links
        FROM documents d
        WHERE (SELECT COUNT(*) FROM document_links WHERE source_doc_id = d.id OR target_doc_id = d.id) > 0
        ORDER BY links DESC
        LIMIT 5
    """).fetchall()

    conn.close()

    print(f"\n{'='*60}")
    print(f"  Knowledge Engine Stats")
    print(f"{'='*60}\n")
    print(f"  Documents:        {stats['documents']}")
    print(f"  Sources:          {stats['sources']}")
    print(f"  Entities:         {entity_count}")
    print(f"  Document links:   {link_count}")
    print(f"  Synonyms:         {synonym_count}")
    print(f"  Collections:      {collection_count}")
    print(f"  Feedback events:  {stats['feedback_events']}")
    print(f"  Total queries:    {stats['queries']}")
    print(f"  Avg relevance:    {stats['avg_relevance']:.3f}")

    if entity_types:
        print(f"\n  Entity Types:")
        for et in entity_types:
            print(f"    {et['type']:15s} {et['cnt']}")

    if connected:
        print(f"\n  Most Connected Documents:")
        for c in connected:
            print(f"    [{c['id']}] {c['title'][:45]} ({c['links']} links)")

    print()


def cmd_ingest(args):
    """Ingest a local text file directly."""
    db.init_db()
    filepath = args.file
    if not os.path.exists(filepath):
        print(f"File not found: {filepath}")
        return

    with open(filepath, "r") as f:
        content = f.read()

    title = args.title or os.path.basename(filepath)
    keywords = scraper.extract_keywords(content)
    h = scraper.content_hash(content)

    conn = db.get_conn()
    source = conn.execute("SELECT id FROM sources WHERE type='manual' AND name='local-files'").fetchone()
    if not source:
        db.add_source("manual", "local-files")
        source = conn.execute("SELECT id FROM sources WHERE type='manual' AND name='local-files'").fetchone()
    conn.close()

    doc_id = db.insert_document(
        source_id=source["id"],
        title=title,
        content=content,
        keywords=keywords,
        doc_hash=h
    )
    if doc_id:
        print(f"Ingested: {title} (id: {doc_id}, {len(keywords)} keywords)")
        # Extract entities automatically
        ent_count = entities.extract_and_store(doc_id, f"{title} {content}")
        if ent_count:
            print(f"  Extracted {ent_count} entities")
    else:
        print(f"Skipped (duplicate): {title}")


def cmd_extract_entities(args):
    """Run entity extraction on all documents."""
    db.init_db()
    print("Extracting entities from all documents...")
    total = entities.process_all_documents()
    print(f"\nDone. Extracted {total} total entity references.")

    if args.auto_link:
        print(f"\nAuto-linking documents with >= {args.min_shared} shared entities...")
        links = entities.auto_link_documents(min_shared=args.min_shared)
        print(f"Created {links} document links.")


def cmd_entity_search(args):
    """Search by entity."""
    db.init_db()
    name = args.name
    etype = args.type

    results = smart_search.entity_search(entity_name=name, entity_type=etype, limit=args.limit)

    if not results:
        print("No documents found for that entity.")
        return

    print(f"\n{'='*60}")
    print(f"  Entity Search: {name or ''} ({etype or 'all types'})")
    print(f"  Found: {len(results)} documents")
    print(f"{'='*60}\n")

    for r in results:
        print(f"  [{r['id']}] {r['title']}")
        print(f"       Entity: {r['entity_name']} ({r['entity_type']}) — mentions: {r['frequency']}")
        snippet = r["content"][:150].replace("\n", " ")
        print(f"       {snippet}...")
        print()


def cmd_entities(args):
    """List top entities."""
    db.init_db()
    results = smart_search.get_top_entities(limit=args.limit, entity_type=args.type)

    if not results:
        print("No entities found. Run 'extract-entities' first.")
        return

    print(f"\n{'='*60}")
    print(f"  Top Entities{' (' + args.type + ')' if args.type else ''}")
    print(f"{'='*60}\n")

    for e in results:
        print(f"  [{e['id']}] {e['name']:35s} ({e['type']:10s}) — in {e['actual_count']} docs")
    print()


def cmd_graph(args):
    """Walk the document relationship graph."""
    db.init_db()
    doc_id = args.id

    conn = db.get_conn()
    doc = conn.execute("SELECT title FROM documents WHERE id = ?", (doc_id,)).fetchone()
    conn.close()

    if not doc:
        print(f"Document {doc_id} not found.")
        return

    results = smart_search.graph_walk(doc_id, depth=args.depth, limit=args.limit)

    if not results:
        print(f"No connected documents found for [{doc_id}] {doc['title']}")
        return

    print(f"\n{'='*60}")
    print(f"  Graph Walk from: [{doc_id}] {doc['title']}")
    print(f"  Depth: {args.depth} | Found: {len(results)} connected docs")
    print(f"{'='*60}\n")

    for r in results:
        indent = "  " * r["depth"]
        print(f"  {indent}[{r['id']}] {r['title']}")
        print(f"  {indent}     ← {r['relation']} (strength: {r['strength']:.2f}, depth: {r['depth']})")
    print()


def cmd_related(args):
    """Find documents related to a given document by shared entities."""
    db.init_db()
    doc_id = args.id
    results = entities.find_related_documents(doc_id, limit=args.limit)

    conn = db.get_conn()
    doc = conn.execute("SELECT title FROM documents WHERE id = ?", (doc_id,)).fetchone()
    conn.close()

    if not results:
        print(f"No related documents found for [{doc_id}] {doc['title'] if doc else '?'}")
        return

    print(f"\n{'='*60}")
    print(f"  Related to: [{doc_id}] {doc['title']}")
    print(f"{'='*60}\n")

    for r in results:
        shared_names = r["shared_entity_names"] or ""
        top_shared = ", ".join(shared_names.split(",")[:5])
        print(f"  [{r['id']}] {r['title']}")
        print(f"       Shared entities: {r['shared_entities']} — {top_shared}")
    print()


def cmd_seed_synonyms(args):
    """Seed the synonym table with domain-specific terms."""
    db.init_db()
    count = smart_search.seed_synonyms()
    print(f"Seeded {count} synonym pairs for query expansion.")


def cmd_add_synonym(args):
    """Add a custom synonym pair."""
    db.init_db()
    smart_search.add_synonym(args.term, args.synonym, args.weight)
    print(f"Added synonym: {args.term} ↔ {args.synonym} (weight: {args.weight})")


def cmd_collection_create(args):
    """Create a new collection."""
    db.init_db()
    conn = db.get_conn()
    conn.execute("INSERT OR IGNORE INTO collections (name, description) VALUES (?, ?)",
                 (args.name, args.description))
    conn.commit()
    conn.close()
    print(f"Created collection: {args.name}")


def cmd_collection_add(args):
    """Add a document to a collection."""
    db.init_db()
    conn = db.get_conn()
    coll = conn.execute("SELECT id FROM collections WHERE name = ?", (args.collection,)).fetchone()
    if not coll:
        print(f"Collection '{args.collection}' not found.")
        conn.close()
        return
    conn.execute("INSERT OR IGNORE INTO collection_documents (collection_id, document_id) VALUES (?, ?)",
                 (coll["id"], args.doc_id))
    conn.commit()
    conn.close()
    print(f"Added document {args.doc_id} to collection '{args.collection}'")


def cmd_collections(args):
    """List all collections."""
    db.init_db()
    conn = db.get_conn()
    rows = conn.execute("""
        SELECT c.*, COUNT(cd.document_id) AS doc_count
        FROM collections c
        LEFT JOIN collection_documents cd ON cd.collection_id = c.id
        GROUP BY c.id
        ORDER BY c.name
    """).fetchall()
    conn.close()

    if not rows:
        print("No collections yet.")
        return

    print(f"\n{'='*60}")
    print(f"  Collections")
    print(f"{'='*60}\n")
    for r in rows:
        print(f"  {r['name']:25s} — {r['doc_count']} docs — {r['description'] or ''}")
    print()


def cmd_scrape_properties(args):
    db.init_db()
    max_price = args.max_price or 30000
    properties_scraper.scrape_all_bc(max_price=max_price)


def cmd_hydrate_properties(args):
    db.init_db()
    properties_scraper.hydrate_kijiji_listings(limit=args.limit or 500)


def cmd_find_property(args):
    db.init_db()
    query = " ".join(args.query) if args.query else None
    rows = db.search_properties(
        query=query,
        max_price=args.max_price or 30000,
        region=args.region,
        ptype=args.type,
        limit=args.limit,
    )
    if not rows:
        print("No properties found.")
        return

    print(f"\n{'='*60}")
    print(f"  Found {len(rows)} properties")
    print(f"{'='*60}\n")
    for r in rows:
        price = f"${r['price']:,}" if r["price"] else "?"
        hidden = " [HIDDEN]" if r["is_hidden"] else ""
        print(f"  [{r['id']}] {r['title']}{hidden}")
        print(f"       Price: {price} | Region: {r['region'] or '?'} | Type: {r['property_type'] or '?'}")
        if r["size_acres"]:
            print(f"       Size: {r['size_acres']} acres")
        if r["location"]:
            print(f"       Location: {r['location']}")
        if r["listing_url"]:
            print(f"       URL: {r['listing_url']}")
        print(f"       Score: {r['score']:.1f} | Source: {r['source'] or '?'}")
        if r["description"]:
            snippet = r["description"][:180].replace("\n", " ")
            print(f"       {snippet}...")
        print()


def cmd_property_stats(args):
    db.init_db()
    s = db.property_stats()
    print(f"\n{'='*60}")
    print(f"  BC Cheap Properties Stats")
    print(f"{'='*60}\n")
    print(f"  Total listings:   {s['total']}")
    print(f"  Under $30k:       {s['under_30k']}")
    print(f"  Under $10k:       {s['under_10k']}")
    print(f"  Hidden/obscure:   {s['hidden']}")
    print(f"  Avg price:        ${s['avg_price']:,.0f}")
    if s["by_region"]:
        print(f"\n  By region:")
        for region, count in sorted(s["by_region"].items(), key=lambda x: -x[1]):
            print(f"    {region or '?':20s} {count}")
    if s["by_type"]:
        print(f"\n  By type:")
        for ptype, count in sorted(s["by_type"].items(), key=lambda x: -x[1]):
            print(f"    {ptype or '?':20s} {count}")
    if s["by_source"]:
        print(f"\n  By source:")
        for src, count in sorted(s["by_source"].items(), key=lambda x: -x[1]):
            print(f"    {src or '?':20s} {count}")
    print()


def cmd_decay(args):
    db.init_db()
    conn = db.get_conn()
    decay_rate = args.rate or 0.01
    conn.execute("""
        UPDATE documents
        SET relevance_score = relevance_score + (0.5 - relevance_score) * ?
        WHERE feedback_count > 0
    """, (decay_rate,))
    affected = conn.execute("SELECT changes()").fetchone()[0]
    conn.commit()
    conn.close()
    print(f"Applied decay (rate={decay_rate}) to {affected} documents.")


def main():
    parser = argparse.ArgumentParser(description="Knowledge Engine — Advanced Self-Refining Knowledge Base")
    sub = parser.add_subparsers(dest="command")

    # init
    sub.add_parser("init", help="Initialize the database")

    # add-source
    p = sub.add_parser("add-source", help="Add a data source")
    p.add_argument("--type", required=True, choices=["rss", "web", "api", "manual"])
    p.add_argument("--name", required=True)
    p.add_argument("--url")
    p.add_argument("--selector", help="CSS selector for web scraping")

    # scrape
    p = sub.add_parser("scrape", help="Scrape all enabled sources")
    p.add_argument("--extract-entities", action="store_true", help="Extract entities after scraping")

    # query (now with smart search)
    p = sub.add_parser("query", help="Search the knowledge base")
    p.add_argument("query", nargs="+")
    p.add_argument("--limit", type=int, default=10)
    p.add_argument("--min-relevance", type=float, default=0.0)
    p.add_argument("--smart", action="store_true", help="Use composite scoring engine")
    p.add_argument("--topic", help="Filter by topic")
    p.add_argument("--no-recency", action="store_true", help="Disable recency boost")
    p.add_argument("--no-expand", action="store_true", help="Disable query expansion")

    # feedback
    p = sub.add_parser("feedback", help="Give relevance feedback")
    p.add_argument("--id", type=int, required=True)
    p.add_argument("--relevant", required=True)
    p.add_argument("--context", help="The query context for this feedback")

    # sources
    sub.add_parser("sources", help="List all sources")

    # stats
    sub.add_parser("stats", help="Show database stats")

    # ingest
    p = sub.add_parser("ingest", help="Ingest a local file (with auto entity extraction)")
    p.add_argument("--file", required=True)
    p.add_argument("--title")

    # decay
    p = sub.add_parser("decay", help="Apply relevance decay")
    p.add_argument("--rate", type=float, default=0.01)

    # === NEW ADVANCED COMMANDS ===

    # extract-entities
    p = sub.add_parser("extract-entities", help="Extract entities from all documents")
    p.add_argument("--auto-link", action="store_true", help="Auto-create document links from shared entities")
    p.add_argument("--min-shared", type=int, default=3, help="Min shared entities to auto-link")

    # entity-search
    p = sub.add_parser("entity-search", help="Search documents by entity")
    p.add_argument("--name", help="Entity name to search for")
    p.add_argument("--type", choices=["person", "place", "org", "concept", "law", "plant", "mineral", "date"])
    p.add_argument("--limit", type=int, default=20)

    # entities
    p = sub.add_parser("entities", help="List top entities")
    p.add_argument("--type", choices=["person", "place", "org", "concept", "law", "plant", "mineral", "date"])
    p.add_argument("--limit", type=int, default=30)

    # graph
    p = sub.add_parser("graph", help="Walk the document relationship graph")
    p.add_argument("--id", type=int, required=True, help="Starting document ID")
    p.add_argument("--depth", type=int, default=2, help="How many hops to traverse")
    p.add_argument("--limit", type=int, default=20)

    # related
    p = sub.add_parser("related", help="Find related documents by shared entities")
    p.add_argument("--id", type=int, required=True)
    p.add_argument("--limit", type=int, default=10)

    # seed-synonyms
    sub.add_parser("seed-synonyms", help="Seed domain-specific synonyms for query expansion")

    # add-synonym
    p = sub.add_parser("add-synonym", help="Add a custom synonym pair")
    p.add_argument("--term", required=True)
    p.add_argument("--synonym", required=True)
    p.add_argument("--weight", type=float, default=0.8)

    # collection-create
    p = sub.add_parser("collection-create", help="Create a document collection")
    p.add_argument("--name", required=True)
    p.add_argument("--description", default="")

    # collection-add
    p = sub.add_parser("collection-add", help="Add a document to a collection")
    p.add_argument("--collection", required=True)
    p.add_argument("--doc-id", type=int, required=True)

    # collections
    sub.add_parser("collections", help="List all collections")

    # property commands
    p = sub.add_parser("scrape-properties", help="Scrape cheap BC properties")
    p.add_argument("--max-price", type=int, default=30000)

    p = sub.add_parser("hydrate-properties", help="Fetch detail pages for listings missing data")
    p.add_argument("--limit", type=int, default=500)

    p = sub.add_parser("find-property", help="Search saved cheap properties")
    p.add_argument("query", nargs="*")
    p.add_argument("--max-price", type=int, default=30000)
    p.add_argument("--region")
    p.add_argument("--type")
    p.add_argument("--limit", type=int, default=20)

    sub.add_parser("property-stats", help="Show property database stats")

    args = parser.parse_args()
    if not args.command:
        parser.print_help()
        return

    commands = {
        "init": cmd_init,
        "add-source": cmd_add_source,
        "scrape": cmd_scrape,
        "query": cmd_query,
        "feedback": cmd_feedback,
        "sources": cmd_sources,
        "stats": cmd_stats,
        "ingest": cmd_ingest,
        "decay": cmd_decay,
        "extract-entities": cmd_extract_entities,
        "entity-search": cmd_entity_search,
        "entities": cmd_entities,
        "graph": cmd_graph,
        "related": cmd_related,
        "seed-synonyms": cmd_seed_synonyms,
        "add-synonym": cmd_add_synonym,
        "collection-create": cmd_collection_create,
        "collection-add": cmd_collection_add,
        "collections": cmd_collections,
        "scrape-properties": cmd_scrape_properties,
        "hydrate-properties": cmd_hydrate_properties,
        "find-property": cmd_find_property,
        "property-stats": cmd_property_stats,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()