import os
import json
import hashlib
import threading
from openai import OpenAI, AsyncOpenAI
from tqdm import tqdm
import sqlite3
import traceback
from tqdm.asyncio import tqdm_asyncio
import asyncio
from concurrent.futures import Future


def gen_task_hash(messages, model_name, params) -> str:
    """根据输入内容生成唯一哈希，用于缓存命中"""
    raw = json.dumps(
        {"messages": messages, "model": model_name, "params": params},
        sort_keys=True,
        ensure_ascii=False
    )
    return hashlib.md5(raw.encode("utf-8")).hexdigest()


class SQLiteCache:
    def __init__(self, path="cache.db"):
        self.conn = sqlite3.connect(path, check_same_thread=False)
        self.conn.execute('PRAGMA journal_mode=WAL;')
        self.conn.execute("""
        CREATE TABLE IF NOT EXISTS cache(
            hash TEXT PRIMARY KEY,
            data TEXT NOT NULL
        )
        """)
        self.conn.commit()

    def get(self, hash_key):
        cur = self.conn.execute("SELECT data FROM cache WHERE hash=?", (hash_key,))
        row = cur.fetchone()
        return json.loads(row[0]) if row else None

    def save(self, hash_key, record):
        self.conn.execute(
            "REPLACE INTO cache(hash, data) VALUES (?, ?)",
            (hash_key, json.dumps(record, ensure_ascii=False))
        )
        self.conn.commit()
    
    def get_many(self, hash_keys):
        """
        批量读取 cache
        :param hash_keys: List[str]
        :return: Dict[str, Any]  {hash: record}
        """
        if not hash_keys:
            return {}

        # SQLite 对 IN (...) 有参数数量限制（默认 999）
        MAX_VARS = 900
        result = {}

        for i in tqdm(range(0, len(hash_keys), MAX_VARS)):
            batch = hash_keys[i:i + MAX_VARS]
            placeholders = ",".join("?" for _ in batch)

            sql = f"SELECT hash, data FROM cache WHERE hash IN ({placeholders})"
            cur = self.conn.execute(sql, batch)
            exe_results = cur.fetchall()
            for h, data in exe_results:
                result[h] = json.loads(data)

        return result

class APIClient:
    def __init__(self, url, api_key, cache_path="/mnt/temp_cache/cache.db", max_try=3, max_concurrent=6, timeout=120):
        self.url = url
        self.api_key = api_key
        self.max_try = max_try
        self.timeout = timeout
        self.cache = SQLiteCache(cache_path)
        self.client = AsyncOpenAI(base_url=self.url, api_key=self.api_key, timeout=self.timeout)
        self.semaphore = asyncio.Semaphore(max_concurrent)

        # ---- 创建后台事件循环线程 ----
        self.loop = asyncio.new_event_loop()
        self.loop_thread = threading.Thread(target=self._loop_runner, daemon=True)
        self.loop_thread.start()
    
    def get_model_list(self):
        client = OpenAI(base_url=self.url, api_key=self.api_key)
        model_list = client.models.list()
        return [model.id for model in model_list.data]

    def _loop_runner(self):
        """后台线程运行事件循环"""
        asyncio.set_event_loop(self.loop)
        self.lock = asyncio.Lock()
        self.loop.run_forever()

    async def _query_api(self, messages, model_name, params):
        async with self.semaphore:
            for i in range(self.max_try):
                try:
                    resp = await self.client.chat.completions.create(
                        model=model_name,
                        messages=messages,
                        **params
                    )
                    response = resp.choices[0].message.content if resp.choices[0].message.content else {"error": "return_empty", "traceback": "resp", "retry_count": "return_empty"}
                    return response
                except Exception as e:
                    tb = traceback.format_exc()
                    print(f"[Retry {i+1}/{self.max_try}] Error on model {model_name}: {e}")
                    print(tb)
                    # await asyncio.sleep(30*2^i) # 指数退避
                    if i + 1 == self.max_try:
                        return {"error": str(e), "traceback": tb, "retry_count": i + 1}

    async def _run_async(self, input_datas, show_progress=True):
        print("begin hash")
        hash_list = [gen_task_hash(inp[0], inp[1], inp[2]) for inp in input_datas]

        tasks_to_run = []
        results = []

        print("search hash sql")
        async with self.lock:
            cached_map = self.cache.get_many(hash_list)  # dict: {hash: record}

        # 2. 本地内存里判断
        for idx, h in enumerate(hash_list):
            cached = cached_map.get(h)
            if cached:
                results.append(cached)
            else:
                results.append(None)
                tasks_to_run.append((idx, input_datas[idx], h))

        print("search hash sql finished")
        async def handle_task(idx, data, h):
            result = await self._query_api(data[0], data[1], data[2])
            if isinstance(result, str):
                record = {
                    "hash": h,
                    "messages": data[0],
                    "model": data[1],
                    "params": data[2],
                    "result": result
                }
                async with self.lock:
                    self.cache.save(h, record)
                return idx, record
            else:
                return idx, {
                    "error": result["error"],
                    "traceback": result["traceback"],
                    "retry_count": result["retry_count"]
                }

        if tasks_to_run:
            tasks = [handle_task(idx, data, h) for idx, data, h in tasks_to_run]

            if show_progress:
                for fut in tqdm_asyncio.as_completed(tasks, total=len(tasks), desc="Processing"):
                    idx, res = await fut
                    results[idx] = res
            else:
                for idx, res in await asyncio.gather(*tasks):
                    results[idx] = res

        return results

    def run(self, input_datas, show_progress=True):
        future: Future = asyncio.run_coroutine_threadsafe(
            self._run_async(input_datas, show_progress),
            self.loop
        )
        return future.result()   # 阻塞等待结果（线程安全）
    
def sqlite_to_jsonl(sqlite_path):
    jsonl_path = sqlite_path.replace(".db", ".jsonl")
    assert sqlite_path!=jsonl_path
    conn = sqlite3.connect(sqlite_path)
    cursor = conn.cursor()

    cursor.execute("SELECT data FROM cache")
    rows = cursor.fetchall()

    with open(jsonl_path, "w", encoding="utf-8") as f:
        for row in rows:
            data_str = row[0]
            f.write(data_str + "\n")

    conn.close()
    print(f"导出完成：{len(rows)} 条记录 → {jsonl_path}")

if __name__ == "__main__":
    cache_path = "/mnt/temp_cache/cache.db"
    client = APIClient(
        url="Your API URL", 
        api_key="Your API Key",   
        cache_path=cache_path,
        max_concurrent=10,
        max_try=3
    )
    import random
    extra_body={"thinking": {"type": "disabled"}}
    test_datas = [
        [[{"role": "user", "content": "Hello, who are you?"}], "doubao-seed-1.6-lite", {"temperature": random.random(), "extra_body": extra_body}],
        [[{"role": "user", "content": "This should fail"}], "gemini-3-pro-preview", {"temperature": 0, "extra_body": extra_body}], 
        [[{"role": "user", "content": "Tell me a joke"}], "deepseek-v3.1", {"temperature": 0, "extra_body": extra_body}],
    ]*2

    results = client.run(test_datas, show_progress=True)
    sqlite_to_jsonl(cache_path)
    a = 1