cli.py

import sys
import os

# Validate API key early
api_key = os.getenv("GROQ_API_KEY", "")
if not api_key:
	print("\n[Error] GROQ_API_KEY environment variable is not set.")
	print("Get your free key at: https://console.groq.com")
	print("Then run: set GROQ_API_KEY=gsk_...")
	sys.exit(1)

from rich.console import Console
from rich.panel   import Panel
from rich.table   import Table
from rich         import print as rprint

try:
	from .graph_store  import GraphStore
	from .extractor    import Extractor
	from .reasoner     import Reasoner
	from .visualizer   import render as render_graph
	from .config       import VISUAL_OUTPUT, DATA_DIR, VERSION
	from .wikipedia    import fetch_article, search_article
except ImportError:
	from graph_store  import GraphStore
	from extractor    import Extractor
	from reasoner     import Reasoner
	from visualizer   import render as render_graph
	from config       import VISUAL_OUTPUT, DATA_DIR, VERSION
	from wikipedia    import fetch_article, search_article

console = Console()

BANNER = f"""
[bold cyan]╔══════════════════════════════════════╗[/bold cyan]
[bold cyan]║    Knowledge Graph Reasoner  v{VERSION}    ║[/bold cyan]
[bold cyan]╚══════════════════════════════════════╝[/bold cyan]
Type [bold yellow]help[/bold yellow] for commands.
"""

def print_help():
	table = Table(show_header=True, header_style="bold magenta", box=None, padding=(0,2))
	table.add_column("Command", style="bold yellow")
	table.add_column("Description")

	commands = [
		("ingest <path>",               "Ingest a .txt file into the graph"),
		("ingest-text",                 "Paste text directly (blank line to end)"),
		("ingest-wiki <title>",         "Fetch and ingest a Wikipedia article by name"),
		("ingest-wiki-batch",           "Ingest multiple Wikipedia articles in one go"),
		("wiki-search <query>",         "Search Wikipedia for article titles"),
		("sources",                     "List all ingested documents"),
		("delete-source <name>", "Remove a source and its exclusive nodes from the graph"),
		("query <question>",            "Ask a question answered from graph data"),
		("lookup <entity>",             "See all relationships for an entity"),
		("connect <A> :: <B>",          "Find the path between two entities"),
		("contradictions",              "Scan the graph for contradictions"),
		("visualize",                   "Render the full graph to HTML"),
		("visualize <entity>",          "Render local subgraph around an entity"),
		("visualize <entity> <depth>",  "Render subgraph with custom depth (e.g. 3)"),
		("stats",                       "Show graph statistics"),
		("search <term>",               "Search for nodes by name"),
		("reset",                       "Wipe the graph (asks confirmation)"),
		("help",                        "Show this help"),
		("quit / exit",                 "Exit"),
	]
	for cmd, desc in commands:
		table.add_row(cmd, desc)

	console.print(Panel(table, title="Commands", border_style="cyan"))

def print_stats(graph_store: GraphStore):
	G = graph_store.graph

	table = Table(show_header=False, box=None, padding=(0,2))
	table.add_column("Key", style="bold cyan")
	table.add_column("Value", style="white")

	table.add_row("Nodes", str(G.number_of_nodes()))
	table.add_row("Edges", str(G.number_of_edges()))

	if G.number_of_nodes() > 0:
		top = sorted(G.degree(), key=lambda x: x[1], reverse=True)[:5]
		top_str = ", ".join(f"{n} ({d})" for n, d in top)
		table.add_row("Top nodes", top_str)

		components = len(list(__import__("networkx").weakly_connected_components(G)))
		table.add_row("Components",	str(components))

	console.print(Panel(table, title="Graph Statistics", border_style="cyan"))

def ingest_pasted_text(extractor: Extractor):
	console.print("[cyan]Paste your text below. Enter a blank line when done:[/cyan]")
	lines = []
	while True:
		try:
			line = input()
			if line == "":
				if lines:
					break
			else:
				lines.append(line)
		except EOFError:
			break

	if not lines:
		console.print("[yellow]No text entered.[/yellow]")
		return

	text = "\n".join(lines)
	name = console.input("[cyan]Source name (e.g. 'wikipedia_article'): [/cyan]").strip()
	if not name:
		name = "pasted_text"

	extractor.ingest_text(text, source_name=name)

def main():
	DATA_DIR.mkdir(exist_ok=True)

	console.print(BANNER)

	# Initialise components
	console.print("[dim]Loading graph...[/dim]")
	graph_store = GraphStore()

	console.print("[dim]Loading models (first run downloads ~90MB)...[/dim]")
	extractor = Extractor(graph_store)
	reasoner  = Reasoner(graph_store)

	console.print(f"[green]Ready.[/green] Graph has {graph_store.node_count()} nodes, {graph_store.edge_count()} edges.\n")

	while True:
		try:
			raw = console.input("[bold cyan]> [/bold cyan]").strip()
		except (EOFError, KeyboardInterrupt):
			console.print("\n[yellow]Goodbye.[/yellow]")
			break

		if not raw:
			continue

		parts = raw.split(None, 1)
		cmd   = parts[0].lower()
		rest  = parts[1].strip() if len(parts) > 1 else ""

		# Commands

		if cmd in ("quit", "exit", "q"):
			console.print("[yellow]Goodbye.[/yellow]")
			break

		elif cmd == "help":
			print_help()

		elif cmd == "stats":
			print_stats(graph_store)

		elif cmd == "ingest":
			if not rest:
				console.print("[red]Usage: ingest <path/to/file.txt>[/red]")
			else:
				try:
					extractor.ingest_file(rest)
				except FileNotFoundError as e:
					console.print(f"[red]{e}[/red]")

		elif cmd == "ingest-text":
			ingest_pasted_text(extractor)
		
		elif cmd == "ingest-wiki":
			if not rest:
				console.print("[red]Usage: ingest-wiki <Article Title>[/red]")
				console.print("[dim]Tip: use wiki-search to find the exact title first[/dim]")
			else:
				try:
					console.print(f"[dim]Fetching '{rest}' from Wikipedia...[/dim]")
					canonical, text = fetch_article(rest)
					console.print(f"[green]Fetched:[/green] {canonical} ({len(text):,} chars)")
					extractor.ingest_text(text, source_name=canonical)
				except (ValueError, RuntimeError) as e:
					console.print(f"[red]{e}[/red]")
		
		elif cmd == "ingest-wiki-batch":
			console.print("[cyan]Enter article titles one per line. Blank line to start ingesting:[/cyan]")
			titles = []
			while True:
				try:
					line = input("  title: ").strip()
					if not line:
						if titles:
							break
					else:
						titles.append(line)
				except EOFError:
					break
			
			if not titles:
				console.print("[yellow]No titles entered.[/yellow]")
			else:
				console.print(f"\n[cyan]Ingesting {len(titles)} article(s)...[/cyan]\n")
				ok, fail = 0, 0
				for i, title in enumerate(titles, 1):
					console.print(f"[dim][{i}/{len(titles)}] Fetching '{title}'...[/dim]")
					try:
						canonical, text = fetch_article(title)
						console.print(f"  [green]✓[/green] {canonical} ({len(text):,} chars)")
						extractor.ingest_text(text, source_name=canonical)
						ok += 1
					except (ValueError, RuntimeError) as e:
						console.print(f"  [red]✗ {e}[/red]")
						fail +=1
				console.print(f"\n[green]Done:[/green] {ok} ingested, {fail} failed.")
		
		elif cmd == "wiki-search":
			if not rest:
				console.print("[red]Usage: wiki-search <query>[/red]")
			else:
				results = search_article(rest)
				if results:
					console.print(f"[green]Wikipedia results for '{rest}':[/green]")
					for r in results:
						console.print(f"  • {r}")
					console.print("[dim]Use the exact title with ingest-wiki[/dim]")
				else:
					console.print(f"[yellow]No results found for '{rest}'.[/yellow]")
		
		elif cmd == "sources":
			sources = graph_store.all_sources()
			if sources:
				console.print(f"[green]{len(sources)} source(s) ingested:[/green]")
				for s in sources:
					console.print(f"  • {s}")
			else:
				console.print("[yellow]No sources ingested yet.[/yellow]")
		
		elif cmd == "delete-source":
			if not rest:
				console.print("[red]Usage: delete-source <source name>[/red]")
				console.print("[dim]Use 'sources' to see exact names[/dim]")
			else:
				all_sources = graph_store.all_sources()
				match = next((s for s in all_sources if s.lower() == rest.lower()), None)
				if not match:
					console.print(f"[red]No source matching '{rest}' found.[/red]")
					console.print("[dim]Use 'sources' to see exact names[/dim]")
				else:
					confirm = console.input(f"[red]Delete source '{match}' and its nodes? (yes/no): [/red]").strip().lower()
					if confirm == "yes":
						result = graph_store.delete_source(match)
						console.print(f"[green]Deleted:[/green] {result['removed_nodes']} nodes removed, {result['kept_nodes']} nodes kept (shared with other sources)")
					else:
						console.print("[yellow]Cancelled.[/yellow]")

		elif cmd == "query":
			if not rest:
				console.print("[red]Usage: query <your question>[/red]")
			else:
				console.print("[dim]Querying graph...[/dim]")
				answer = reasoner.query(rest)
				console.print(Panel(answer, title="Answer", border_style="green"))

		elif cmd == "lookup":
			if not rest:
				console.print("[red]Usage: lookup <entity name>[/red]")
			else:
				result = reasoner.lookup_entity(rest)
				console.print(Panel(result, title=f"Lookup: {rest}", border_style="blue"))

		elif cmd == "connect":
			# Expects format: connect Entity A :: Entity B
			if "::" not in rest:
				console.print("[red]Usage: connect <Entity A> :: <Entity B>[/red]")
			else:
				a, b = [x.strip() for x in rest.split("::", 1)]
				console.print(f"[dim]Finding path between '{a}' and '{b}'...[/dim]")
				result = reasoner.find_connection(a, b)
				console.print(Panel(result, title=f"Connection: {a} ↔ {b}", border_style="magenta"))

		elif cmd == "contradictions":
			console.print("[dim]Scanning for contradictions...[/dim]")
			result = reasoner.find_contradictions()
			console.print(Panel(result, title="Contradictions", border_style="red"))

		elif cmd == "search":
			if not rest:
				console.print("[red]Usage: search <term>[/red]")
			else:
				matches = graph_store.search_nodes(rest)
				if matches:
					console.print(f"[green]Found {len(matches)} node(s):[/green]")
					for m in matches[:20]:
						console.print(f"  • {m}")
				else:
					console.print(f"[yellow]No nodes found matching '{rest}'.[/yellow]")

		elif cmd == "visualize":
			focus = None
			depth = 2
			if rest:
				parts2 = rest.rsplit(None, 1)
				if len(parts2) == 2 and parts2[1].isdigit():
					focus = parts2[0].strip()
					depth = int(parts2[1])
				else:
					focus = rest.strip()
			
			console.print("[dim]Rendering graph...[/dim]")
			try:
				path = render_graph(graph_store, focus_node=focus, depth=depth)
				console.print(f"[green]Graph rendered to:[/green] {path}")
				console.print("[dim]Open that file in your browser.[/dim]")
			except ValueError as e:
				console.print(f"[red]{e}[/red]")
			except Exception as e:
				console.print(f"[red]Visualization failed: {e}[/red]")

		elif cmd == "reset":
			confirm = console.input("[red]Are you sure you want to wipe the graph? (yes/no): [/red]").strip().lower()
			if confirm == "yes":
				graph_store.reset()
			else:
				console.print("[yellow]Cancelled.[/yellow]")

		else:
			console.print(f"[red]Unknown command: '{cmd}'. Type 'help' for commands.[/red]")

if __name__ == "__main__":
	main()