import chromadb
import gradio as gr
from sentence_transformers import SentenceTransformer
from llama_cpp import Llama

# ✅ Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path="./chromadb_store")
collection = chroma_client.get_or_create_collection(name="curly_strings_knowledge")

# ✅ Load Local Embedding Model
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# ✅ Load Fine-Tuned LLaMA Model
llm = Llama.from_pretrained(
    repo_id="krishna195/second_guff",
    filename="unsloth.Q4_K_M.gguf",
)

# ✅ File-Based Search Function
def search_in_file(query, file_path="merged_output.txt"):
    try:
        with open(file_path, "r", encoding="utf-8") as file:
            lines = file.readlines()
        
        # Search for the query in file content
        matched_lines = [line.strip() for line in lines if query.lower() in line.lower()]
        
        return "\n".join(matched_lines) if matched_lines else "No relevant data found in file."
    
    except FileNotFoundError:
        return "File not found. Please check the file path."

# ✅ Retrieve Context from ChromaDB & File
def retrieve_context(query):
    query_embedding = embedder.encode(query).tolist()
    results = collection.query(query_embeddings=[query_embedding], n_results=2)

    retrieved_texts = [doc for sublist in results.get("documents", []) for doc in sublist if isinstance(doc, str)]
    
    # If no result from ChromaDB, try searching in the file
    if not retrieved_texts:
        return search_in_file(query)
    
    return "\n".join(retrieved_texts)

# ✅ Chatbot Function with Optimized Retrieval
def chatbot_response(user_input):
    context = retrieve_context(user_input)

    messages = [
        {"role": "system", "content": """You are an expert on the Estonian folk band Curly Strings.  
        - Use the **retrieved knowledge** from ChromaDB or the file to answer.  
        - If a **song** is mentioned, provide its name and **suggest similar tracks**.  
        - If no match is found, say "I couldn’t find details, but here’s what I know."."""},
        {"role": "user", "content": user_input},
        {"role": "assistant", "content": f"Retrieved Context:\n{context}"},
    ]

    response = llm.create_chat_completion(
        messages=messages,
        temperature=0.4,  
        max_tokens=300,
        top_p=0.9,
        frequency_penalty=0.7,
    )

    return response["choices"][0]["message"]["content"].strip()

# ✅ Gradio Chatbot Interface
iface = gr.Interface(
    fn=chatbot_response,
    inputs=gr.Textbox(label="Ask me about Curly Strings 🎻"),
    outputs=gr.Textbox(label="Bot Response 🎶"),
    title="Curly Strings Chatbot",
    description="Ask about the Estonian folk band Curly Strings! Now also searches in 'merged_output.txt'.",
)

iface.launch()