"""
rapid-rag CLI - Command line interface for local RAG.
"""

import argparse
import sys
from pathlib import Path


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        prog="rapid-rag",
        description="Fast local RAG - search your documents with AI"
    )
    subparsers = parser.add_subparsers(dest="command", help="Commands")

    # Init command
    init_parser = subparsers.add_parser("init", help="Initialize a new RAG collection")
    init_parser.add_argument("name", help="Collection name")
    init_parser.add_argument("--dir", help="Persist directory", default=None)

    # Add command
    add_parser = subparsers.add_parser("add", help="Add documents")
    add_parser.add_argument("path", help="File or directory to add")
    add_parser.add_argument("-c", "--collection", default="default", help="Collection name")
    add_parser.add_argument("-r", "--recursive", action="store_true", help="Recursive directory scan")
    add_parser.add_argument("--chunk-size", type=int, default=1000, help="Chunk size")
    add_parser.add_argument("--chunk-overlap", type=int, default=200, help="Chunk overlap")

    # Search command
    search_parser = subparsers.add_parser("search", help="Semantic search")
    search_parser.add_argument("query", help="Search query")
    search_parser.add_argument("-c", "--collection", default="default", help="Collection name")
    search_parser.add_argument("-n", "--num", type=int, default=5, help="Number of results")

    # Query command (RAG with LLM)
    query_parser = subparsers.add_parser("query", help="RAG query with LLM")
    query_parser.add_argument("question", help="Question to answer")
    query_parser.add_argument("-c", "--collection", default="default", help="Collection name")
    query_parser.add_argument("-m", "--model", default="qwen2.5:7b", help="Ollama model")
    query_parser.add_argument("-n", "--num", type=int, default=5, help="Context documents")

    # Info command
    info_parser = subparsers.add_parser("info", help="Collection info")
    info_parser.add_argument("-c", "--collection", default="default", help="Collection name")

    # Clear command
    clear_parser = subparsers.add_parser("clear", help="Clear collection")
    clear_parser.add_argument("-c", "--collection", default="default", help="Collection name")
    clear_parser.add_argument("-y", "--yes", action="store_true", help="Skip confirmation")

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        return 0

    # Import here to avoid slow startup
    from .core import RapidRAG

    if args.command == "init":
        rag = RapidRAG(args.name, persist_dir=args.dir)
        print(f"Initialized collection '{args.name}'")
        print(f"Persist dir: {rag.persist_dir}")
        return 0

    if args.command == "add":
        rag = RapidRAG(args.collection)
        path = Path(args.path)

        if path.is_file():
            ids = rag.add_file(
                path,
                chunk_size=args.chunk_size,
                chunk_overlap=args.chunk_overlap
            )
            print(f"Added {len(ids)} chunks from {path.name}")
        elif path.is_dir():
            ids = rag.add_directory(
                path,
                recursive=args.recursive,
                chunk_size=args.chunk_size,
                chunk_overlap=args.chunk_overlap
            )
            print(f"Added {len(ids)} chunks total")
        else:
            print(f"Path not found: {path}", file=sys.stderr)
            return 1

        return 0

    if args.command == "search":
        rag = RapidRAG(args.collection)
        results = rag.search(args.query, n_results=args.num)

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

        for i, r in enumerate(results, 1):
            score = r["score"]
            source = r["metadata"].get("source", r["id"])
            content = r["content"][:200].replace("\n", " ")
            print(f"\n[{i}] {source} (score: {score:.3f})")
            print(f"    {content}...")

        return 0

    if args.command == "query":
        rag = RapidRAG(args.collection)
        result = rag.query(args.question, n_results=args.num, model=args.model)

        print("\n" + "=" * 60)
        print("ANSWER:")
        print("=" * 60)
        print(result["answer"])
        print("\n" + "-" * 60)
        print("SOURCES:")
        for s in result["sources"]:
            source = s["metadata"].get("source", s["id"])
            print(f"  - {source} (score: {s['score']:.3f})")

        return 0

    if args.command == "info":
        rag = RapidRAG(args.collection)
        print(f"Collection: {args.collection}")
        print(f"Documents: {rag.count()}")
        print(f"Persist dir: {rag.persist_dir}")
        print(f"Embedding model: {rag.embedding_model}")
        return 0

    if args.command == "clear":
        if not args.yes:
            confirm = input(f"Clear collection '{args.collection}'? [y/N] ")
            if confirm.lower() != "y":
                print("Cancelled.")
                return 0

        rag = RapidRAG(args.collection)
        rag.clear()
        print(f"Cleared collection '{args.collection}'")
        return 0

    return 0


if __name__ == "__main__":
    sys.exit(main())
