cli.py

"""
Codebase Assistant — Interactive CLI
Run with: .venv\Scripts\python.exe cli.py
"""

import os
import sys
import asyncio

# ── styling ───────────────────────────────────────────────────────────────────

CYAN   = "\033[96m"
GREEN  = "\033[92m"
YELLOW = "\033[93m"
RED    = "\033[91m"
BOLD   = "\033[1m"
DIM    = "\033[2m"
RESET  = "\033[0m"

# ── session state ─────────────────────────────────────────────────────────────
# Tracks the repo path provided at startup (or via re-ingest) so that option 7
# can save generated files directly into that directory without asking again.
_current_repo: str = "."

def banner():
    print(f"""
{CYAN}{BOLD}
  ╔══════════════════════════════════════════╗
  ║       AI  Codebase  Assistant  CLI       ║
  ╚══════════════════════════════════════════╝
{RESET}""")

def section(title: str):
    print(f"\n{CYAN}{BOLD}── {title} {'─' * (42 - len(title))}{RESET}")

def ok(msg: str):
    print(f"  {GREEN}✓{RESET}  {msg}")

def err(msg: str):
    print(f"  {RED}✗  {msg}{RESET}")

def info(msg: str):
    print(f"  {DIM}{msg}{RESET}")

def ask(prompt: str) -> str:
    return input(f"\n{YELLOW}▶ {prompt}{RESET} ").strip()

# ── ingest ────────────────────────────────────────────────────────────────────

def ingest_repository(root_path: str):
    global _current_repo
    section("Ingesting Repository")
    from rag.repository_loader import load_repository
    from rag.splitter import split_code
    from rag.embedding import embed_document
    from rag.retriever import build_vector_store

    print(f"  {DIM}Loading files from: {root_path}{RESET}")
    docs = load_repository(root_path)
    if not docs:
        err("No supported files found. Check allowed_extensions in config.py.")
        return False
    ok(f"Loaded {len(docs)} files")

    all_chunks = []
    for doc in docs:
        all_chunks.extend(split_code(doc))
    ok(f"Split into {len(all_chunks)} chunks")

    print(f"  {DIM}Embedding chunks (this may take a moment)...{RESET}")
    embedded = embed_document(all_chunks)
    ok(f"Embedded {len(embedded)} chunks")

    build_vector_store(embedded)
    ok("Vector store ready (ChromaDB)\n")
    _current_repo = root_path   # remember for option 7
    return True

# ── feature handlers ──────────────────────────────────────────────────────────

def feature_query():
    section("Ask a Question  [RAG Query]")
    query = ask("Your question:")
    if not query:
        return
    k_str = ask("Number of source chunks to retrieve? [default: 5]")
    k = int(k_str) if k_str.isdigit() else 5

    from rag.rag_chain import run_rag_query
    print(f"\n  {DIM}Thinking...{RESET}")
    try:
        result = run_rag_query(query, k=k)
        section("Answer")
        print(f"\n{result['answer']}")
        section("Sources")
        for s in result["sources"]:
            print(f"  {DIM}• {s.get('file_path')}  lines {s.get('start_line')}–{s.get('end_line')}{RESET}")
    except Exception as e:
        err(str(e))


def feature_bugs():
    section("Detect Bugs  [LLM Code Review]")
    file_path = ask("File path to analyse:")
    if not os.path.isfile(file_path):
        err(f"File not found: {file_path}")
        return

    from services.code_analysis import detect_bugs
    print(f"\n  {DIM}Analysing...{RESET}")
    try:
        bugs = detect_bugs(file_path)
        if not bugs:
            ok("No issues found.")
            return
        section(f"Found {len(bugs)} issue(s)")
        for b in bugs:
            sev = b.get("severity", "?").upper()
            color = RED if sev in ("HIGH", "CRITICAL") else YELLOW
            print(f"\n  {color}[{sev}]{RESET} Line {b.get('line', '?')}: {b.get('issue', '')}")
            print(f"  {DIM}  → {b.get('suggestion', '')}{RESET}")
    except Exception as e:
        err(str(e))


def feature_complexity():
    section("Cyclomatic Complexity Analysis")
    file_path = ask("File path to analyse:")
    if not os.path.isfile(file_path):
        err(f"File not found: {file_path}")
        return

    from services.code_analysis import analyze_complexity
    try:
        result = analyze_complexity(file_path)
        fns = result["functions"]
        if not fns:
            info("No functions found (or file is not Python).")
            return
        section(f"{len(fns)} function(s)")
        for fn in fns:
            rank = fn["rank"]
            color = GREEN if rank == "A" else (YELLOW if rank in ("B", "C") else RED)
            bar = "█" * fn["complexity"]
            print(f"  {color}[{rank}]{RESET}  {fn['name']:<30} complexity={fn['complexity']}  {DIM}{bar}{RESET}")
    except Exception as e:
        err(str(e))


def feature_explain():
    section("Explain a Function")
    file_path = ask("File path:")
    if not os.path.isfile(file_path):
        err(f"File not found: {file_path}")
        return
    func_name = ask("Function name:")
    if not func_name:
        return

    from services.code_analysis import explain_function
    print(f"\n  {DIM}Thinking...{RESET}")
    try:
        explanation = explain_function(file_path, func_name)
        section("Explanation")
        print(f"\n{explanation}")
    except Exception as e:
        err(str(e))


def feature_module_docs():
    section("Generate Module Documentation")
    file_path = ask("File path:")
    if not os.path.isfile(file_path):
        err(f"File not found: {file_path}")
        return

    from services.documentation import generate_module_docs
    print(f"\n  {DIM}Generating docs...{RESET}")
    try:
        docs = generate_module_docs(file_path)
        section("Documentation")
        print(f"\n{docs}")
    except Exception as e:
        err(str(e))


def feature_readme():
    section("Generate README.md")
    root = ask("Repository root path:")
    if not os.path.isdir(root):
        err(f"Directory not found: {root}")
        return

    from services.documentation import generate_readme
    print(f"\n  {DIM}Generating README...{RESET}")
    try:
        readme = generate_readme(root)
        section("README Preview")
        print(f"\n{readme[:1500]}{'...' if len(readme) > 1500 else ''}")
        save = ask("Save to README.md in that directory? [y/N]")
        if save.lower() == "y":
            out = os.path.join(root, "README.md")
            with open(out, "w", encoding="utf-8") as f:
                f.write(readme)
            ok(f"Saved to {out}")
    except Exception as e:
        err(str(e))


def feature_propose_file():
    section("Propose a New File  [AI File Creator]")
    description = ask("Describe the file you want to create:")
    if not description:
        return
    context_query = ask("Optional context query (or press Enter to skip):")

    from services.file_creator import propose_new_file
    print(f"\n  {DIM}Generating...{RESET}")
    try:
        proposal = propose_new_file(description, context_query or None)
        section("Proposed File")
        print(f"\n  {BOLD}Path:{RESET}  {proposal.path}")
        print(f"\n{DIM}{'─'*50}{RESET}")
        print(proposal.content[:1000] + ("..." if len(proposal.content) > 1000 else ""))
        print(f"{DIM}{'─'*50}{RESET}")

        confirm = ask("Write this file to disk? [y/N]")
        if confirm.lower() == "y":
            default_dir = _current_repo
            typed = ask(f"Save under which directory? [default: {default_dir}]").strip()
            base_dir = typed if typed else default_dir
            if not os.path.isdir(base_dir):
                err(f"Directory not found: {base_dir}")
                return
            async def _write():
                from services.file_creator import apply_approved_file
                proposal.approved = True
                return await apply_approved_file(proposal, True, base_dir=base_dir)
            result = asyncio.run(_write())
            if result.get("status") == "written":
                ok(f"File written: {result.get('path', proposal.path)}")
            else:
                err(f"Write failed: {result}")
    except Exception as e:
        err(str(e))


def feature_mcp_tools():
    section("List GitHub MCP Tools")
    async def _list():
        from mcp_client import get_github_mcp_client
        async with get_github_mcp_client() as client:
            return await client.list_tools()
    try:
        tools = asyncio.run(_list())
        ok(f"{len(tools)} tools available:")
        for t in tools:
            print(f"  {DIM}• {t['name']:<30}{RESET} {t.get('description','')[:60]}")
    except Exception as e:
        err(str(e))


def feature_reingest():
    root = ask("New repository path to ingest:")
    if os.path.isdir(root):
        ingest_repository(root)
    else:
        err(f"Directory not found: {root}")

# ── menu ──────────────────────────────────────────────────────────────────────

MENU = [
    ("Ask a question about the codebase",    feature_query),
    ("Detect bugs in a file",                feature_bugs),
    ("Cyclomatic complexity analysis",       feature_complexity),
    ("Explain a function",                   feature_explain),
    ("Generate module documentation",        feature_module_docs),
    ("Generate README for a repository",     feature_readme),
    ("Propose & create a new file (AI)",     feature_propose_file),
    ("List GitHub MCP tools",                feature_mcp_tools),
    ("Re-ingest a repository",               feature_reingest),
]

def show_menu():
    section("What would you like to do?")
    for i, (label, _) in enumerate(MENU, 1):
        print(f"  {CYAN}{i:>2}{RESET}  {label}")
    print(f"  {DIM}  0  Exit{RESET}")

# ── main ──────────────────────────────────────────────────────────────────────

def main():
    # Enable ANSI colours on Windows
    os.system("")

    banner()

    # Step 1 — ask for repo path
    while True:
        root = ask("Enter the path to the repository you want to analyse:")
        if os.path.isdir(root):
            break
        err(f"Directory not found: {root}")

    # Step 2 — ingest
    if not ingest_repository(root):
        sys.exit(1)

    # Step 3 — interactive menu loop
    while True:
        show_menu()
        choice = ask("Choose an option [0–9]:")
        if choice == "0" or choice.lower() in ("exit", "quit", "q"):
            print(f"\n{DIM}Goodbye!{RESET}\n")
            break
        if choice.isdigit() and 1 <= int(choice) <= len(MENU):
            MENU[int(choice) - 1][1]()
        else:
            err("Invalid choice. Enter a number from the menu.")


if __name__ == "__main__":
    main()