import json
import uuid
import sqlite3
from typing import Any, Dict, Optional

class OffloadManager:
    """
    負責管理 Tool 之間傳遞的大數據 (Fat Data)。
    資料儲存在 SQLite 中，避免污染 LLM Context。
    """
    DB_PATH = "storage/offloaded_data.db"

    def __init__(self):
        self._init_db()

    def _init_db(self):
        import os
        os.makedirs("storage", exist_ok=True)
        with sqlite3.connect(self.DB_PATH) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS data_store (
                    ref_id TEXT PRIMARY KEY,
                    data JSON
                )
            """)

    def store(self, data: Any) -> str:
        """將數據存入 DB，返回引用 ID"""
        ref_id = f"ref_{uuid.uuid4().hex[:8]}"
        with sqlite3.connect(self.DB_PATH) as conn:
            conn.execute(
                "INSERT INTO data_store (ref_id, data) VALUES (?, ?)",
                (ref_id, json.dumps(data, default=str)) # Handle datetime serialization if needed
            )
        return ref_id

    def retrieve(self, ref_id: str) -> Optional[Any]:
        """根據引用 ID 取回數據"""
        with sqlite3.connect(self.DB_PATH) as conn:
            cursor = conn.execute("SELECT data FROM data_store WHERE ref_id = ?", (ref_id,))
            row = cursor.fetchone()
            if row:
                return json.loads(row[0])
        return None

# Global Instance
offload_manager = OffloadManager()