#from utils.es import ElasticsearchOperations
import asyncio
import os
import time
from utils.file import load_file
from libs.text_processor import TextProcessor
import json
import jieba
import hashlib
import numpy as np
import codecs
from dotenv import load_dotenv

load_dotenv()

class ChuncHelper:    
    def __init__(self, data_file: str, output_dir:str, user_dict:str) -> None:
        self.data_file = data_file
        #self.es = ElasticsearchOperations()
        self.processor = TextProcessor()
        self.output_dir = output_dir
        self.stop_words = set()
        if user_dict:
            jieba.load_userdict(user_dict)
            with open(os.getenv("JIEBA_STOP_DICT"),"r",encoding="utf-8") as f:
                for line in f:
                    self.stop_words.add(line.strip())
                    
        self.json_data = None
        self.load_data()
    def get_doc_id(self, title:str) -> str:
        md = hashlib.md5(title.encode())
        doc_id = md.hexdigest()
        return doc_id
    
    def get_url(self, title:str)->str:
        if self.json_data:
            title = title.strip()
            for item in self.json_data:                
                ArticleTitle = item["meta_data"]["ArticleTitle"]
                if title == ArticleTitle:
                    return item["url"]
        return "未找到相关内容"
        
    def get(self, title:str)->str:
        if self.json_data:
            title = title.strip()
            for item in self.json_data:                
                ArticleTitle = item["meta_data"]["ArticleTitle"]
                if title == ArticleTitle:
                    return item["article_text"]
        return "未找到相关内容"
    
    def get_article(self, title:str):
        result = {}
        if self.json_data:
            title = title.strip()
            for item in self.json_data:                
                ArticleTitle = item["meta_data"]["ArticleTitle"]
                if title == ArticleTitle:
                    result["site_name"] = item["meta_data"]["SiteName"]
                    result["site_domain"] = item["meta_data"]["SiteDomain"]
                    result["title"] = item["meta_data"]["ArticleTitle"]
                    result["author"] = item["meta_data"]["author"]
                    result["pub_date"] = item["meta_data"]["PubDate"]
                    result["article_text"] = item["article_text"]
                    result["url"] = item["url"]
                    return result
        return None
    def load_data(self):        
        with open(self.data_file,"r",encoding="utf-8") as f:
            lines = f.readlines()
            self.json_data = json.loads(''.join(lines))  
    def code_t(self, text):
        t = codecs.encode(text, "utf-8")
        return codecs.decode(t, "utf-8")
    
    def cut_word(self, text: str):
        cut_result = jieba.cut(text)
        filtered_words = [word for word in cut_result if word not in self.stop_words]
        return filtered_words
    def title_reverse_index(self):
        ''' 对文章标题分词，然后基于词建立倒排索引 '''
        if self.json_data:
            index = 1
            words_vector = {}
            total = len(self.json_data)
            for item in self.json_data:
                print(f"\r {index}/{total}", end="")
                title = item["meta_data"]["ArticleTitle"]
                title_cut = jieba.cut(title)
                for word in title_cut:
                    if word in words_vector.keys():
                        if not title in words_vector[word]["articles"]:
                            words_vector[word]["articles"].append(title)
                        continue
                    words_vector[word] = { "articles": [title] }
                index = index + 1     
            with open(f"{self.output_dir}/title_reverse_index.json", "w", encoding="utf-8") as f:
                f.write(json.dumps(words_vector,ensure_ascii=False))
            
    def cut_title_vector(self, seperate=False):
        '''基于文章标题分词，每个词生成向量'''
        words_vector = json.loads(load_file(f"{self.output_dir}/words_vector.json"))
        if self.json_data:
            index = 1
            total = len(self.json_data)
            for item in self.json_data:
                print(f"\r {index}/{total}", end="")
                title = item["meta_data"]["ArticleTitle"]
                title_cut = jieba.cut(title)
                for word in title_cut:
                    if word in words_vector.keys():
                        continue
                    words_vector[word] = self.processor.generate_embeddings([word])[0].tolist()
                index = index + 1
        print("\nwriting words vector files")
        
        if seperate == False:
            with open(f"{self.output_dir}/words_vector.json", "w", encoding="utf-8") as f:
                f.write(json.dumps(cached_json,ensure_ascii=False))
            return
        
        cached_json = {}
        count = 0
        size = 100
        index = 0
        for k in words_vector.keys():
            count += 1
            cached_json[k] = words_vector[k]
            if count == size:                              
                with open(f"{self.output_dir}/words_vector_{index}.json", "w", encoding="utf-8") as f:
                    f.write(json.dumps(cached_json,ensure_ascii=False))
                cached_json = {}
                count = 0
                index += 1
        index += 1
        if len(cached_json)>0:
            with open(f"{self.output_dir}/words_vector_{index}.json", "w", encoding="utf-8") as f:
                f.write(json.dumps(cached_json,ensure_ascii=False))
    def process_data(self): 
        '''对每个文章进行处理，生成关键词向量，标题向量，内容分片向量'''          
        def default_converter(o):
            if isinstance(o, np.float32):
                return float(o)
            raise TypeError            
        
        if self.json_data:
            total = len(self.json_data)
            print(f">>> total {total} documents ")
            index = 1
            data_collection = []
            for item in self.json_data:                
                print(f">>> process {index}/{total}")
                data = {"title":self.code_t(item["meta_data"]['ArticleTitle']),
                        "keywords":self.code_t(item["meta_data"]['keywords'])}
                            
                author = item["meta_data"]["author"]
                description = item["meta_data"]["description"]
                keywords = item["meta_data"]["keywords"]
                SiteName = item["meta_data"]["SiteName"]
                SiteDomain = item["meta_data"]["SiteDomain"]
                SiteIDCode = item["meta_data"]["SiteIDCode"]
                ColumnName = item["meta_data"]["ColumnName"]
                ColumnType = item["meta_data"]["ColumnType"]
                ArticleTitle = item["meta_data"]["ArticleTitle"]
                doc_id = self.get_doc_id(ArticleTitle)
                
                if os.path.exists(f"{self.output_dir}/{doc_id}.json"):
                    print(f"{doc_id} existed, skip")
                    index = index + 1
                    continue
                PubDate = item["meta_data"]["PubDate"]
                ContentSource = item["meta_data"]["ContentSource"]
                article_text = item["article_text"]
                text_len = len(article_text)
                print(f"{ArticleTitle}:{text_len}")
                #txt = [self.code_t(keywords),self.code_t(ArticleTitle),self.code_t(article_text)]                
                txt = [self.code_t(keywords),self.code_t(ArticleTitle)]                
                chuncs = self.processor.chunk_text(self.code_t(article_text))
                txt = txt + chuncs
                print(">>> start embedding...")
                embeded_text = self.processor.generate_embeddings(txt)
                #embeded_chuncs = self.processor.generate_embeddings(chuncs)
                title_cut = jieba.cut(ArticleTitle)
                keywords_cut = jieba.cut(keywords)
                data["title"] = ArticleTitle
                data["title_cut"] = list(title_cut)
                data["keywords_cut"] = list(keywords_cut)
                data["keywords_vector"] = embeded_text[0].tolist()
                data["title_vector"] = embeded_text[1].tolist()
                data['chuncs'] = chuncs
                #data["content_vector"] = embeded_text[2].tolist()
                data["chuncs_vector"] = embeded_text[2:].tolist()
                
                print(">>> write embedding...")
                with open(f"{self.output_dir}/{doc_id}.json", "w", encoding="utf-8") as f:
                    f.write(json.dumps(data,ensure_ascii=False))
                
                print(f"{doc_id} done, {index}/{total}")
                index = index + 1
   