from datetime import datetime
import os
import json
from typing import Any, Dict, List, Optional, Union
import torch


class ModelExistsError(Exception):
    """exception raised when existing model found with the same configuration at a certain path"""

    def __init__(self, message="Model already exists!"):
        self.message = message
        super().__init__(self.message)


class ModelNotFoundError(Exception):
    """exception raised when model with desired configuration not found at a certain path"""

    def __init__(self, message="No matching model found!"):
        self.message = message
        super().__init__(self.message)


class Storage():
    """Simple storage class for saving and loading models using JSON metadata"""

    def __init__(self, base_directory: str, experiment_name: Optional[str] = ''):
        cache_dir = os.path.join(base_directory, experiment_name)
        os.makedirs(cache_dir, exist_ok=True)
        self.cache_dir = cache_dir

    def _get_metadata_path(self, artifact_type: str) -> str:
        return os.path.join(self.cache_dir, f'{artifact_type}_metadata.json')

    def _load_metadata(self, artifact_type: str) -> List[Dict]:
        metadata_path = self._get_metadata_path(artifact_type)
        if not os.path.exists(metadata_path):
            return []
        
        try:
            with open(metadata_path, 'r') as f:
                return json.load(f)
        except:
            return []

    def _save_metadata(self, artifact_type: str, metadata: List[Dict]) -> None:
        metadata_path = self._get_metadata_path(artifact_type)
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, default=str)

    def build_model_dir_path(self, artifact_type: str, dir_id: Union[int, str]) -> str:
        path = os.path.join(self.cache_dir, artifact_type, str(dir_id))
        os.makedirs(path, exist_ok=True)
        return path

    def build_model_file_path(self, path: str, init_no: int = 1) -> str:
        model_file_name = f'model_{init_no}.pt'
        model_file_path = os.path.join(path, model_file_name)
        return model_file_path

    def create_model_file_path(self, artifact_type: str, params: Dict[str, Any], init_no: int = 1) -> str:
        metadata = self._load_metadata(artifact_type)
        
        # Check for matching params
        matching_entries = [entry for entry in metadata if entry['params'] == params]
        
        # If there's already a matching entry, use it
        if matching_entries:
            entry_id = matching_entries[0]['id']
            path = self.build_model_dir_path(artifact_type, entry_id)
            model_file_path = self.build_model_file_path(path, init_no)
            
            if os.path.exists(model_file_path):
                raise ModelExistsError(f'Model already exists (artifact_type={artifact_type}, params={params}, init_no={init_no})')
            
            return model_file_path
        
        # Otherwise create a new entry
        new_id = len(metadata) + 1
        entry = {
            'id': new_id,
            'params': params,
            'time': datetime.utcnow().isoformat(),
        }
        
        metadata.append(entry)
        self._save_metadata(artifact_type, metadata)
        
        path = self.build_model_dir_path(artifact_type, new_id)
        model_file_path = self.build_model_file_path(path, init_no)
        
        return model_file_path

    def retrieve_model_dir_path(self, artifact_type: str, match_condition: Dict[str, Any]) -> str:
        metadata = self._load_metadata(artifact_type)
        
        # Find entries that match all conditions
        matching_entries = []
        for entry in metadata:
            matches = True
            for key, value in match_condition.items():
                if key not in entry['params'] or entry['params'][key] != value:
                    matches = False
                    break
            if matches:
                matching_entries.append(entry)
        
        if not matching_entries:
            raise ModelNotFoundError(f'No matching model found with (artifact_type={artifact_type}, match_condition={match_condition})')
        
        # Use the first matching entry
        document = matching_entries[0]
        return self.build_model_dir_path(artifact_type, document['id'])

    def retrieve_model_file_path(self, artifact_type: str, params: Dict[str, Any], init_no: int = 1) -> str:
        path = self.retrieve_model_dir_path(artifact_type, params)
        model_file_path = self.build_model_file_path(path, init_no)
        
        if not os.path.exists(model_file_path):
            raise ModelNotFoundError(f'No matching model found with (artifact_type={artifact_type}, match_condition={params})')
        
        return model_file_path

    def find_artifacts(self, artifact_type: str, match_condition: Dict[str, Any], exact_params: bool = False) -> List[Dict[str, Any]]:
        metadata = self._load_metadata(artifact_type)
        documents = []
        
        for entry in metadata:
            if exact_params:
                if entry['params'] == match_condition:
                    documents.append(entry)
            else:
                matches = True
                for key, value in match_condition.items():
                    if key not in entry['params'] or entry['params'][key] != value:
                        matches = False
                        break
                if matches:
                    documents.append(entry)
        
        return documents
