import json
import os
import random
from typing import Union, Dict, Any
from cryptography.fernet import Fernet

class CognitionCocooner:
    def __init__(self, storage_path: str = "cocoons", encryption_key: bytes = None):
        self.storage_path = storage_path
        os.makedirs(self.storage_path, exist_ok=True)
        self.key = encryption_key or Fernet.generate_key()
        self.fernet = Fernet(self.key)

    def wrap(self, thought: Dict[str, Any], type_: str = "prompt") -> str:
        cocoon = {
            "type": type_,
            "id": f"cocoon_{random.randint(1000,9999)}",
            "wrapped": self._generate_wrapper(thought, type_)
        }
        file_path = os.path.join(self.storage_path, cocoon["id"] + ".json")

        with open(file_path, "w") as f:
            json.dump(cocoon, f)

        return cocoon["id"]

    def unwrap(self, cocoon_id: str) -> Union[str, Dict[str, Any]]:
        file_path = os.path.join(self.storage_path, cocoon_id + ".json")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Cocoon {cocoon_id} not found.")

        with open(file_path, "r") as f:
            cocoon = json.load(f)

        return cocoon["wrapped"]

    def wrap_encrypted(self, thought: Dict[str, Any]) -> str:
        encrypted = self.fernet.encrypt(json.dumps(thought).encode()).decode()
        cocoon = {
            "type": "encrypted",
            "id": f"cocoon_{random.randint(10000,99999)}",
            "wrapped": encrypted
        }
        file_path = os.path.join(self.storage_path, cocoon["id"] + ".json")

        with open(file_path, "w") as f:
            json.dump(cocoon, f)

        return cocoon["id"]

    def unwrap_encrypted(self, cocoon_id: str) -> Dict[str, Any]:
        file_path = os.path.join(self.storage_path, cocoon_id + ".json")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Cocoon {cocoon_id} not found.")

        with open(file_path, "r") as f:
            cocoon = json.load(f)

        decrypted = self.fernet.decrypt(cocoon["wrapped"].encode()).decode()
        return json.loads(decrypted)

    def _generate_wrapper(self, thought: Dict[str, Any], type_: str) -> Union[str, Dict[str, Any]]:
        if type_ == "prompt":
            return f"What does this mean in context? {thought}"
        elif type_ == "function":
            return f"def analyze(): return {thought}"
        elif type_ == "symbolic":
            return {k: round(v, 2) for k, v in thought.items()}
        else:
            return thought