"""
Forensic Logger - Blockchain-style Chain of Custody
Implements tamper-evident logging with cryptographic hashing
"""

import json
import hashlib
import datetime
import uuid
from pathlib import Path
from typing import Dict, Any, Optional
import logging


class ForensicLogger:
    """
    Forensic logging with blockchain-style chain of custody
    Each event is hashed with the previous event's hash
    """
    
    def __init__(self, log_dir: str = "logs/forensic", experiment_id: Optional[str] = None):
        """Initialize forensic logger"""
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        # generate unique ID if not provided
        self.experiment_id = experiment_id or self._generate_experiment_id()
        
        # chain tracking
        self.event_chain = []
        self.previous_hash = None  # starts as None for first event
        
        # setup logger
        self.logger = self._setup_logger()
        
        # log the initialization itself
        self._log_initialization()
    
    def _generate_experiment_id(self):
        """Generate unique experiment ID with timestamp"""
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        uid = str(uuid.uuid4())[:8]
        return f"EXP_{timestamp}_{uid}"
    
    def _setup_logger(self):
        """Setup JSON line logger"""
        log_file = self.log_dir / f"{self.experiment_id}.jsonl"
        
        logger = logging.getLogger(f"ForensicLogger_{self.experiment_id}")
        logger.setLevel(logging.INFO)
        logger.handlers.clear()  # clear any existing handlers
        
        handler = logging.FileHandler(log_file)
        handler.setLevel(logging.INFO)
        logger.addHandler(handler)
        
        return logger
    
    def _compute_event_hash(self, event_data):
        """
        Compute SHA-256 hash including previous hash
        This creates the blockchain-style chain
        """
        # get previous hash (None for first event)
        prev_hash = event_data.get("previous_hash", "")
        if prev_hash is None:
            prev_hash = ""
        
        # serialize event data and combine with previous hash
        data_string = json.dumps(event_data, sort_keys=True)
        combined = f"{prev_hash}{data_string}"
        
        return hashlib.sha256(combined.encode()).hexdigest()
    
    def _log_initialization(self):
        """Log the logger initialization itself"""
        self.log_event(
            event_type="SYSTEM_INIT",
            data={
                "experiment_id": self.experiment_id,
                "log_directory": str(self.log_dir)
            },
            message="Forensic logging system initialized"
        )
    
    def log_event(self, event_type, data, message=""):
        """
        Log an event with cryptographic chain of custody
        
        Returns event_id for reference
        """
        event_id = str(uuid.uuid4())
        
        # build event structure
        event = {
            "event_id": event_id,
            "event_type": event_type,
            "timestamp": datetime.datetime.now().isoformat(),
            "experiment_id": self.experiment_id,
            "data": data,
            "message": message,
            "previous_hash": self.previous_hash  # link to previous event
        }
        
        # compute hash for this event
        event_hash = self._compute_event_hash(event)
        event["hash"] = event_hash
        
        # update chain state
        self.event_chain.append(event_id)
        self.previous_hash = event_hash  # next event will use this
        
        # write to JSONL file
        self.logger.info(json.dumps(event))
        
        return event_id
    
    def log_data_operation(self, operation, dataset, num_samples, file_hash=None):
        """Log data loading or preprocessing"""
        return self.log_event(
            event_type="DATA_OPERATION",
            data={
                "operation": operation,
                "dataset": dataset,
                "num_samples": num_samples,
                "file_hash": file_hash
            },
            message=f"Data operation: {operation} on {dataset}"
        )
    
    def log_model_operation(self, operation, model_name, model_hash=None, parameters=None):
        """Log model operations (load, train, etc)"""
        return self.log_event(
            event_type="MODEL_OPERATION",
            data={
                "operation": operation,
                "model_name": model_name,
                "model_hash": model_hash,
                "parameters": parameters or {}
            },
            message=f"Model operation: {operation} on {model_name}"
        )
    
    def log_inference(self, model_name, input_hash, output, confidence=None):
        """Log individual inference with input/output tracking"""
        return self.log_event(
            event_type="INFERENCE",
            data={
                "model_name": model_name,
                "input_hash": input_hash,
                "output": str(output),
                "confidence": confidence
            },
            message=f"Inference performed by {model_name}"
        )
    
    def log_explanation(self, method, model_name, input_hash, explanation_file):
        """Log XAI explanation generation"""
        return self.log_event(
            event_type="EXPLANATION",
            data={
                "method": method,
                "model_name": model_name,
                "input_hash": input_hash,
                "explanation_file": explanation_file
            },
            message=f"Explanation generated using {method}"
        )
    
    def log_adversarial_test(self, attack_type, epsilon, success_rate=None):
        """Log adversarial attack testing"""
        return self.log_event(
            event_type="ADVERSARIAL_TEST",
            data={
                "attack_type": attack_type,
                "epsilon": epsilon,
                "success_rate": success_rate
            },
            message=f"Adversarial test: {attack_type} with epsilon={epsilon}"
        )
    
    def compute_file_hash(self, file_path):
        """Compute SHA-256 hash of a file"""
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            # read in chunks to handle large files
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    
    def verify_chain_integrity(self):
        """
        Verify the entire chain by recomputing all hashes
        Returns True if chain is valid, False if tampering detected
        """
        log_file = self.log_dir / f"{self.experiment_id}.jsonl"
        
        if not log_file.exists():
            print("ERROR: Log file not found")
            return False
        
        # load all events
        events = []
        with open(log_file, 'r') as f:
            for line in f:
                events.append(json.loads(line))
        
        # verify each hash in the chain
        for i, event in enumerate(events):
            stored_hash = event["hash"]
            
            # recreate event without hash field
            event_without_hash = {k: v for k, v in event.items() if k != "hash"}
            
            # recompute the hash
            recomputed_hash = self._compute_event_hash(event_without_hash)
            
            # check if they match
            if recomputed_hash != stored_hash:
                print(f"ERROR: Integrity violation at event {i+1}/{len(events)}")
                print(f"  Event ID: {event['event_id']}")
                return False
        
        print(f"SUCCESS: Chain integrity verified - {len(events)} events")
        return True
    
    def generate_audit_report(self, output_file=None):
        """Generate comprehensive audit report"""
        log_file = self.log_dir / f"{self.experiment_id}.jsonl"
        
        # load all events
        events = []
        with open(log_file, 'r') as f:
            for line in f:
                events.append(json.loads(line))
        
        # compile statistics
        report = {
            "experiment_id": self.experiment_id,
            "total_events": len(events),
            "event_types": {},
            "timeline": {
                "start": events[0]["timestamp"] if events else None,
                "end": events[-1]["timestamp"] if events else None
            },
            "events": events
        }
        
        # count event types
        for event in events:
            event_type = event["event_type"]
            report["event_types"][event_type] = report["event_types"].get(event_type, 0) + 1
        
        # save if requested
        if output_file:
            output_path = Path(output_file)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(output_path, 'w') as f:
                json.dump(report, f, indent=2)
        
        return report


# quick test
if __name__ == "__main__":
    import os
    
    # cleanup old test logs
    test_log_dir = Path("logs/forensic")
    if test_log_dir.exists():
        for f in test_log_dir.glob("TEST_*.jsonl"):
            f.unlink()
    
    print("Testing Forensic Logger...")
    logger = ForensicLogger(experiment_id="TEST_001")
    
    # test some operations
    logger.log_data_operation(
        operation="LOAD",
        dataset="FairFace",
        num_samples=1000,
        file_hash="abc123def456"
    )
    
    logger.log_model_operation(
        operation="LOAD",
        model_name="FaceNet",
        model_hash="model_hash_xyz",
        parameters={"embedding_size": 512}
    )
    
    # verify integrity
    is_valid = logger.verify_chain_integrity()
    
    # generate report
    report = logger.generate_audit_report()
    print(f"Total events logged: {report['total_events']}")
    
    print("Test completed successfully" if is_valid else "Test failed")