#!/usr/bin/env python3

import os
import sys
import argparse
from code_rag.rag import RAG

def stream_output(response_iter):
    """Stream the response to console, handling each token"""
    try:
        for chunk in response_iter:
            print(chunk, end="", flush=True)
        print()  # New line at end
    except KeyboardInterrupt:
        print("\nStreaming interrupted by user")


def interactive_chat(rag_pipeline, rag_chain, no_stream=False):
    """Run an interactive chat session"""
    print(
        "\nEnter your questions about the codebase. Type 'exit', 'quit', or press Ctrl+C to end the session.\n"
    )

    try:
        while True:
            try:
                query = input("\nQuestion: ").strip()

                if query.lower() in ["exit", "quit", ""]:
                    print("\nEnding chat session.")
                    break

                print("\nResponse:")
                if no_stream:
                    response = rag_pipeline.query_rag(rag_chain, query, stream=False)
                    print(response)
                else:
                    response_iter = rag_pipeline.query_rag(
                        rag_chain, query, stream=True
                    )
                    stream_output(response_iter)

            except KeyboardInterrupt:
                print("\n\nEnding chat session.")
                break
            except EOFError:
                break
    except KeyboardInterrupt:
        print("\nChat session interrupted.")


def main():
    parser = argparse.ArgumentParser(
        description="Code RAG - Query your codebase using natural language"
    )
    parser.add_argument(
        "docs_dir", help="Directory containing the documents to process", type=str
    )
    parser.add_argument(
        "query",
        help="Initial query about your codebase (optional in interactive mode)",
        nargs="?",
        default=None,
        type=str,
    )
    parser.add_argument(
        "--db-dir",
        help="Directory to store the vector database (default: .code_rag_db)",
        default=os.path.expanduser("~/.code_rag_db"),
        type=str,
    )
    parser.add_argument(
        "--tracker-file",
        help="File to track document changes (default: .code_rag_tracker.json)",
        default=os.path.expanduser("~/.code_rag_tracker.json"),
        type=str,
    )
    parser.add_argument(
        "--force-refresh",
        help="Force refresh of the vector database",
        action="store_true",
    )
    parser.add_argument(
        "--ollama-url",
        help="URL for the Ollama server (default: 127.0.0.1)",
        default="127.0.0.1",
        type=str,
    )
    parser.add_argument(
        "--embedding-model",
        help="Model to use for embeddings (default: nomic-embed-text)",
        default="nomic-embed-text",
        type=str,
    )
    parser.add_argument(
        "--llm-model",
        help="Model to use for text generation (default: llama3.2)",
        default="llama3.2",
        type=str,
    )
    parser.add_argument(
        "--no-stream", help="Disable streaming output", action="store_true"
    )
    parser.add_argument(
        "--no-interactive",
        help="Run in non-interactive mode (answer single query and exit)",
        action="store_true",
    )

    args = parser.parse_args()

    if args.no_interactive and not args.query:
        parser.error("Query is required in non-interactive mode")

    # Create RAG pipeline
    rag_pipeline = RAG(
        docs_dir=args.docs_dir,
        db_dir=args.db_dir,
        tracker_file=args.tracker_file,
        ollama_url=args.ollama_url,
        embedding_model=args.embedding_model,
        llm_model=args.llm_model,
    )

    # Create or update vector database
    print(f"Processing documents in {args.docs_dir}...")
    print(f"Using models: embedding={args.embedding_model}, llm={args.llm_model}")
    print(f"Ollama server: {args.ollama_url}")
    vectorstore = rag_pipeline.create_vector_db(force_refresh=args.force_refresh)

    # Set up RAG chain
    print("Setting up RAG pipeline...")
    rag_chain = rag_pipeline.setup_rag()

    if args.no_interactive:
        # Single query mode
        print(f"\nQuery: {args.query}")
        print("\nResponse:")
        if args.no_stream:
            response = rag_pipeline.query_rag(rag_chain, args.query, stream=False)
            stream_output(response)
        else:
            response_iter = rag_pipeline.query_rag(rag_chain, args.query, stream=True)
            stream_output(response_iter)
    else:
        # Interactive mode
        if args.query:
            # Handle initial query if provided
            print(f"\nQuery: {args.query}")
            print("\nResponse:")
            if args.no_stream:
                response = rag_pipeline.query_rag(rag_chain, args.query, stream=False)
                stream_output(response)
            else:
                response_iter = rag_pipeline.query_rag(
                    rag_chain, args.query, stream=True
                )
                stream_output(response_iter)

        # Start interactive chat
        interactive_chat(rag_pipeline, rag_chain, args.no_stream)


if __name__ == "__main__":
    main()
