"""
Forensic Chain Integrity Verification
Verifies the blockchain-style chain of custody in forensic logs
"""

import json
import hashlib
from pathlib import Path
from typing import List, Dict, Tuple
from datetime import datetime


class ForensicChainVerifier:
    """Verify integrity of forensic event chains"""
    
    def __init__(self, log_file):
        self.log_file = Path(log_file)
        self.events = []
        self.load_events()
    
    def load_events(self):
        """Load events from JSONL file"""
        if not self.log_file.exists():
            raise FileNotFoundError(f"Log file not found: {self.log_file}")
        
        with open(self.log_file, 'r') as f:
            for line in f:
                if line.strip():
                    self.events.append(json.loads(line))
    
    def compute_hash(self, event):
        """
        Recompute hash using same algorithm as ForensicLogger
        Need to match exactly or verification will fail
        """
        prev_hash = event.get("previous_hash", "")
        if prev_hash is None:
            prev_hash = ""
        
        # remove hash field before computing
        event_copy = event.copy()
        event_copy.pop("hash", None)
        
        data_string = json.dumps(event_copy, sort_keys=True)
        combined = f"{prev_hash}{data_string}"
        return hashlib.sha256(combined.encode()).hexdigest()
    
    def verify_chain(self):
        """Verify complete chain integrity"""
        if not self.events:
            return False, ["No events found"]
        
        errors = []
        
        # first event should have null previous_hash
        first_event = self.events[0]
        if first_event.get("previous_hash") is not None:
            errors.append(f"First event should have null previous_hash")
        
        # verify each hash
        for i, event in enumerate(self.events):
            stored_hash = event.get("hash")
            if not stored_hash:
                errors.append(f"Event {i} missing hash")
                continue
            
            computed_hash = self.compute_hash(event)
            if stored_hash != computed_hash:
                errors.append(f"Event {i} hash mismatch - possible tampering")
        
        # verify chain linkage
        for i in range(1, len(self.events)):
            prev_event = self.events[i - 1]
            curr_event = self.events[i]
            
            expected_prev_hash = prev_event.get("hash")
            actual_prev_hash = curr_event.get("previous_hash")
            
            if expected_prev_hash != actual_prev_hash:
                errors.append(f"Chain break between events {i-1} and {i}")
        
        return len(errors) == 0, errors
    
    def generate_report(self):
        """Generate verification report"""
        is_valid, errors = self.verify_chain()
        
        timestamps = [e.get("timestamp") for e in self.events if "timestamp" in e]
        
        # count event types
        event_types = {}
        for event in self.events:
            etype = event.get("event_type", "UNKNOWN")
            event_types[etype] = event_types.get(etype, 0) + 1
        
        report = {
            "verification_timestamp": datetime.now().isoformat(),
            "log_file": str(self.log_file),
            "total_events": len(self.events),
            "chain_valid": is_valid,
            "errors": errors if errors else [],
            "timeline": {
                "start": timestamps[0] if timestamps else None,
                "end": timestamps[-1] if timestamps else None
            },
            "event_types": event_types,
            "first_event_hash": self.events[0].get("hash", "")[:16] + "..." if self.events else "",
            "last_event_hash": self.events[-1].get("hash", "")[:16] + "..." if self.events else ""
        }
        
        return report


def verify_all_logs():
    """Verify all forensic logs in directory"""
    script_dir = Path(__file__).parent
    logs_path = script_dir.parent / "logs" / "forensic"
    
    print("="*80)
    print("FORENSIC CHAIN INTEGRITY VERIFICATION")
    print("="*80)
    print(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Log Directory: {logs_path.absolute()}")
    print()
    
    log_files = sorted(logs_path.glob("*.jsonl"))
    
    if not log_files:
        print("No log files found!")
        return
    
    print(f"Found {len(log_files)} log files")
    print("-"*80)
    
    all_valid = True
    total_events = 0
    reports = []
    
    for log_file in log_files:
        print(f"\nVerifying: {log_file.name}")
        
        try:
            verifier = ForensicChainVerifier(log_file)
            report = verifier.generate_report()
            reports.append(report)
            
            total_events += report["total_events"]
            
            if report["chain_valid"]:
                print(f"  ✓ VALID - {report['total_events']} events verified")
            else:
                print(f"  ✗ INVALID - {len(report['errors'])} errors")
                for error in report["errors"][:3]:  # show first 3 errors
                    print(f"     {error}")
                all_valid = False
        
        except Exception as e:
            print(f"  ERROR: {e}")
            all_valid = False
    
    print()
    print("="*80)
    print("VERIFICATION SUMMARY")
    print("="*80)
    print(f"Log Files: {len(log_files)}")
    print(f"Events Verified: {total_events:,}")
    print(f"Status: {'ALL VALID' if all_valid else 'SOME INVALID'}")
    print()
    
    if all_valid:
        print("✓ FORENSIC INTEGRITY CONFIRMED")
        print("  All chains cryptographically verified")
        print("  No tampering detected")
        print("  Chain of custody intact")
    else:
        print("✗ INTEGRITY VIOLATIONS DETECTED")
        print("  Some chains compromised")
    
    print("="*80)
    
    # save report
    output_path = script_dir.parent / "results" / "forensic_reports" / "chain_verification_report.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    full_report = {
        "verification_timestamp": datetime.now().isoformat(),
        "total_log_files": len(log_files),
        "total_events": total_events,
        "all_chains_valid": all_valid,
        "individual_reports": reports
    }
    
    with open(output_path, 'w') as f:
        json.dump(full_report, f, indent=2)
    
    print(f"\nDetailed report: {output_path}")
    
    return all_valid


if __name__ == "__main__":
    verify_all_logs()