import os
import dotenv
from ollama import Client
from langchain_community.vectorstores import FAISS
from langchain_ollama import OllamaEmbeddings

dotenv.load_dotenv()


def load_vector_store():
    embeddings = OllamaEmbeddings(
        model=os.getenv("EMBEDDING_MODEL"),
        base_url=os.getenv("OLLAMA_HOST"),
    )
    return FAISS.load_local(
        "vector_store", embeddings, allow_dangerous_deserialization=True
    )


def ollama_client():
    return Client(
        host=os.getenv("OLLAMA_HOST"),
    )


def ollama_chat(client: Client, prompt: str):
    db = load_vector_store()

    # Get configurable parameters from environment
    k_similar = int(
        os.getenv("SIMILARITY_K", "5")
    )  # Number of similar chunks to retrieve

    # Search for relevant context
    relevant_chunks = db.similarity_search(prompt, k=k_similar)

    context = "\n\n".join([doc.page_content for doc in relevant_chunks])

    return client.chat(
        model=os.getenv("OLLAMA_MODEL"),
        messages=[
            {
                "role": "system",
                "content": """You are a command-line tool expert. Your role is to suggest the most appropriate command-line tools based on the user's needs.

                Rules:
                1. ONLY suggest tools that are explicitly mentioned in the provided context
                2. Start with the most fundamental/common tool for the task
                3. Format your response as:
                   Primary tool: [tool name] - [brief description]
                   Usage: [basic usage example]
                   
                   Related tools:
                   - [alternative tool] - [why/when to use it]
                
                4. If no suitable tool is found in the context, say "I don't see a suitable tool in my current knowledge base for that specific task"
                5. Be concise and direct
                6. Include specific, relevant details from the man pages
                7. Don't make assumptions about tools that aren't in the context

                Here is the relevant documentation context:
                {}
                """.format(
                    context
                ),
            },
            {"role": "user", "content": prompt},
        ],
    )