#import recurive textsplitter
from sentence_transformers import SentenceTransformer
import chromadb
import uuid
import os
from dotenv import load_dotenv
load_dotenv()
from openai import OpenAI

class RAG_1177:
    def __init__(self):
        self.db_name = "RAG_1177"
        
        self.model = SentenceTransformer('KBLab/sentence-bert-swedish-cased')
        self.client = chromadb.PersistentClient(path="RAG_1177_db")
        self.db = self.client.get_or_create_collection(self.db_name)

    
    def get_embeddings(self, text):
        embeddings = self.model.encode(text)
        return (embeddings.tolist())
    
    
    def get_ids(self, num_ids):
        ids = [str(uuid.uuid4()) for _ in range(num_ids)]
        return ids
        
    
    def delete_collection(self):
        self.client.delete_collection(self.db_name)
        return
    
    def retrieve(self, query, num_results):
        query_emb = self.get_embeddings(query)
        result = self.db.query(query_embeddings=query_emb, n_results=num_results, include=['documents', 'metadatas'])
        result_urls = result['metadatas'][0]
        result_docs = result['documents'][0]

        url_list = set([item['url'] for item in result_urls])
        result_urls = "Läs mer på:\n"
        for i, url in enumerate(url_list, start=1):
            result_urls += f"{i}: {url}\n"

        return result_docs, result_urls
    
    def insert(self,docs, emb, urls, ids):
        self.db.add(documents=docs, embeddings=emb, metadatas=urls, ids=ids)
        return 


    def relevant_question(self, query):

        history = []
        system_prompt = """Du ska bedömma frågor och svara endast JA eller NEJ.

        Om frågan är relaterad till 1177 sjukrådgivningsdatabas, hälsa,vård, välmående, sjukrådgivning ska du svara "JA".

        Om frågan inte är relaterad till ovanämnde ämnen ska du svara NEJ"""

        user_prompt = f"""Svara endast JA eller NEJ ifall följande fråga är är relevant till 1177, sjukrådgivning, hälsa, vård och välmående

        Frågan är: {query}"""

            
        history.append({"role": "system", "content": system_prompt})
        history.append({"role": "user", "content": user_prompt})


        client = OpenAI()
        response = client.chat.completions.create(
                model='gpt-3.5-turbo',
                messages=history,
                temperature=0.4,
                stream=False
            )
        return (response.choices[0].message.content)