import sqlite3
from pathlib import Path
import threading
import json
import base64
import pickle

thread_local = threading.local()

class SqliteDB:
    def __init__(self, work_path: str, db_name: str):
        self.work_path = Path(work_path)
        self.db_name = db_name
        self.db_path = self.work_path / f"{db_name}.db"

        self._init_db()

    def _get_conn(self):
        if not hasattr(thread_local, "conn"):
            thread_local.conn = sqlite3.connect(
                self.db_path,
                check_same_thread=False
            )
            thread_local.conn.row_factory = sqlite3.Row
        return thread_local.conn


    def _init_db(self):
        conn = self._get_conn()
        cursor = conn.cursor()

        cursor.execute("""
            CREATE TABLE IF NOT EXISTS records (
                key TEXT PRIMARY KEY,
                meta TEXT,
                output TEXT
            )
        """)
        conn.commit()

    def encoder(self, data):
        """优先使用 JSON，失败则用 pickle，并添加前缀标识"""
        try:
            json_str = json.dumps(data, ensure_ascii=False)
            return "json:" + json_str
        except (TypeError, ValueError):
            serialized = pickle.dumps(data)
            encoded = base64.b64encode(serialized).decode('utf-8')
            return "pickle:" + encoded

    def decoder(self, encoded_data):
        """根据前缀决定使用 JSON 还是 pickle 解码"""
        if encoded_data.startswith("json:"):
            json_str = encoded_data[5:]  # 去掉 "json:" 前缀
            return json.loads(json_str)
        elif encoded_data.startswith("pickle:"):
            b64_str = encoded_data[7:]  # 去掉 "pickle:" 前缀
            decoded_bytes = base64.b64decode(b64_str.encode('utf-8'))
            return pickle.loads(decoded_bytes)

    def set(self, key: str, value: dict):
        conn = self._get_conn()
        cursor = conn.cursor()

        key_str = json.dumps(key,ensure_ascii=False,sort_keys=True)
        meta_str = json.dumps(value["meta"],ensure_ascii=False)
        output_str = self.encoder(value["output"])

        new_value = (key_str, meta_str, output_str)
        SQL= """
            INSERT OR REPLACE INTO records
            (key, meta, output)
            VALUES (?, ?, ?)
        """
        cursor.execute(SQL, new_value)
        conn.commit()

    def get_all(self):
        conn = self._get_conn()
        cursor = conn.cursor()

        SQL= """
            SELECT key, meta, output FROM records
        """
        cursor.execute(SQL)
        res = cursor.fetchall()

        data = {}
        for row in res:
            data[json.loads(row[0])] = {
                "meta": json.loads(row[1]),
                "output": self.decoder(row[2])
            }
        return data