from typing import Optional, Dict, Any, Union, List
from datetime import datetime
from dataclasses import dataclass, field
from uuid import uuid4, UUID
from pymongo import MongoClient
import logging
from .common_utility import AuditAction, AuditStatus, EntityType
from .appconfig_handler import Configurations
import json

@dataclass
class AuditData:
    entity: EntityType
    action: AuditAction
    status: AuditStatus
    message: Optional[str]
    payload: Dict[str, Any]
    timestamp: datetime
    is_validated: bool = True
    parent_id: Optional[str] = None
    other_details: Optional[Dict[str, Any]] = None
    source: str = 'N/A'
    retry_count: int = 1
    unique_id: UUID = field(default_factory=uuid4)
    invocation_id: Optional[str] = None
    tenant_id: Optional[UUID] = None
    automation_id: Optional[int] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "unique_id": self.unique_id,
            "tenant_id": self.tenant_id,
            "invocation_id": self.invocation_id,
            "automation_id": self.automation_id,
            "entity": self.entity.value,
            "action": self.action.value,
            "status": self.status.value,
            "message": self.message,
            "payload": self.payload,
            "timestamp": self.timestamp,
            "other_details": self.other_details,
            "source": self.source,
            "retry_count": self.retry_count,
            "parent_id": self.parent_id,
            "is_validated": self.is_validated
        }

class AuditManager:
    def __init__(self, invocation_id, automation_id, tenant_id, config: Configurations) -> None:
        self.invocation_id = invocation_id
        self.automation_id = automation_id
        self.tenant_id = tenant_id

        self.config: Configurations = json.loads(config.mongo_db_connection_string)
        connection_string = self.config.get("connection_string")
        self.database_name = self.config.get("database_name")
        self.collection_name = self.config.get("collection_name")

        self.client = MongoClient(connection_string, uuidRepresentation="standard")

    def build_audit(self, audit_entries: List[AuditData]) -> None:
        for entry in audit_entries:
            if not entry.invocation_id:
                entry.invocation_id = self.invocation_id
            if not entry.automation_id:
                entry.automation_id = self.automation_id
            if not entry.tenant_id and self.tenant_id:
                entry.tenant_id = UUID(self.tenant_id) if not isinstance(self.tenant_id, UUID) else self.tenant_id

    def audit_validation(self, audit_entries: Union[AuditData, List[AuditData]]) -> bool:
        if isinstance(audit_entries, AuditData):
            return True
        elif isinstance(audit_entries, list) and all(isinstance(e, AuditData) for e in audit_entries):
            return False
        else:
            raise ValueError("Input must be an AuditData instance or a list of AuditData instances.")

    async def save_audit(self, audit_entries: Union[AuditData, List[AuditData]]):
        try:
            is_single = self.audit_validation(audit_entries)
            audit_entries = [audit_entries] if is_single else audit_entries
            self.build_audit(audit_entries)

            db = self.client[self.database_name]
            collection = db[self.collection_name]

            documents = [entry.to_dict() for entry in audit_entries]

            if is_single:
                result = collection.insert_one(documents[0])
                return result.inserted_id
            else:
                result = collection.insert_many(documents)
                return result.inserted_ids

        except Exception as e:
            logging.error(f"An error occurred in saving audit: {e}")
            raise
        finally:
            self.client.close()