from tinydb import TinyDB, Query
import os
import json
from collections import Counter
import gelbooru

class DB:
    def __init__(self):
        db_path = "datas"
        if not os.path.exists(db_path):
            os.makedirs(os.path.dirname(db_path), exist_ok=True)
            json.dump({}, open(db_path+"/db.json", "w", encoding="utf-8"))
        self.db = TinyDB(db_path+"/db.json")
        self.File = Query()
        
    def add(self, data_add):
        self.db.insert(data_add)
    
    def add_raw(self, data_add):
        is_video = False
        is_gif = False
        if data_add["image"].split(".")[-1] in ["webm", "mp4", "avi"]:
            is_video = True
        if data_add["image"].split(".")[-1] == "gif":
            is_gif = True
        add_data = {
            "md5": data_add["md5"],
            "local_filename": data_add["image"],
            "is_video": is_video,
            "is_gif": is_gif,
            "tags": data_add["tags"].split(" "),
            "rating": data_add["rating"],
            "score": data_add["score"],
            "timestamp": data_add["created_at"],
            "display_date": gelbooru.gelbooru_date_parse(data_add["created_at"]),
            "source": data_add["source"],
            "width": data_add["width"],
            "height": data_add["height"]
        }
        self.db.insert(add_data)
        pass
    
    def get_all(self):
        return self.db.all()
    
    def get_search(self, ratings, tags):
        ratings
        results = self.db.search(
            self.File.rating.one_of(self.ratings) &
            self.File.tags.any(self.tags)
        )
        return results
    
    def get_unknow(self):
        results = self.db.search(self.File.tags == [])
        return results
    
    def get_autocomplete(self, tag_slise):
        counter = Counter()
        for record in self.db.all():
            counter.update(record.get("tags", []))
        results = [
            {
                "tag": tag,
                "count": count
            }
            for tag, count in counter.items()
            if tag.lower().startswith(prefix)
        ]
        results.sort(key=lambda x: x["count"], reverse=True)
        top_10 = results[:10]
        return top_10
    
    def edit_data(self, id_file, edit_data):
        self.db.update(self.edit_data, self.File.md5 == self.id_file)