import os
import json
import time
import numpy as np
import pandas as pd
from ollama import Client



class Ragby:

    def __init__(self, chat_model, embedding_model, top_n = 10, ollama_client = None):
        self.chat_model = chat_model
        self.embedding_model = embedding_model
        self.top_n = top_n
        self.ollama_client = Client(host = "http://localhost:11434") if ollama_client is None else ollama_client


    def __repr__(self):
        return f"Ragby(chat_model={self.chat_model}, embedding_model={self.embedding_model}, top_n={self.top_n}, ollama_client={self.ollama_client})"


    def make_chunks(self, data_file_path, paragraphs = False):
        if "chunks" not in os.listdir():
            print("--- No 'chunks' folder found, creating one")
            os.mkdir("chunks")
        file_name = data_file_path.split(".")[0].split("/")[1]
        if paragraphs:
            with open(data_file_path, "r") as f:
                chunks = []
                buffer = []
                for line in f.readlines():
                    line = line.strip()
                    if line:
                        buffer.append(line)
                    elif len(buffer):
                        chunks.append(" ".join(buffer))
                        buffer = []
                if len(buffer):
                    chunks.append(" ".join(buffer))
        else:
            with open(data_file_path, "r") as f:
                chunks = f.read().splitlines()
        with open(f"chunks/{file_name}.json", "w") as f:
            json.dump(chunks, f)
            print(f"--- Chunks stored in 'chunks/{file_name}.json'")
            return [chunks, f"chunks/{file_name}.json"]


    def make_embeddings(self, c):
        if isinstance(c, list):
            return [self.ollama_client.embeddings(model = self.embedding_model, prompt = chunk)["embedding"] for chunk in c]
        elif isinstance(c, str):
            return self.ollama_client.embeddings(model = self.embedding_model, prompt = c)["embedding"]
        else:
            raise ValueError("Input to 'c' must be a string (str) or a list of strings (list[str])")


    def cosine_similarity(self, x, y_arr):
        x_norm = np.linalg.norm(x)
        return [np.dot(x, y) / (x_norm * np.linalg.norm(y)) for y in y_arr]


    def get_embedding_similarity(self, sim_score_list, text_chunks):
        sorted_sim_scores = sorted(zip(sim_score_list, range(len(sim_score_list))), reverse = True)
        return [text_chunks[sss[1]] for sss in sorted_sim_scores[:self.top_n]]


    def make_chat_message(self, user_prompt, system_prompt, text_chunks_path, text_embeddings_path):
        with open(text_embeddings_path, "r") as f:
            text_embeddings = json.load(f)
        with open(text_chunks_path, "r") as f:
            text_chunks = json.load(f)
        prompt_embedding = self.ollama_client.embeddings(model = self.embedding_model, prompt = user_prompt)["embedding"]
        sim_scores = self.cosine_similarity(prompt_embedding, text_embeddings)
        similar_chunks = self.get_embedding_similarity(sim_scores, text_chunks)
        system_content = system_prompt + "\n".join(similar_chunks)
        return [
            {
                "role": "system",
                "content": system_content
            },
            {
                "role": "user",
                "content": user_prompt
            }
        ]


    def chat(self, user_prompt, system_prompt, text_chunks_path, text_embeddings_path):
        chat_messages = self.make_chat_message(
            user_prompt = user_prompt,
            system_prompt = system_prompt,
            text_chunks_path = text_chunks_path,
            text_embeddings_path = text_embeddings_path
        )
        response = self.ollama_client.chat(
            model = self.chat_model,
            messages = chat_messages
        )
        return response["message"]["content"]
