import os
import hashlib
from agent_node.config import SYNC_DIR
from protos import agent_pb2

class NodeSyncManager:
    """Handles local filesystem synchronization on the Agent Node."""
    def __init__(self, base_sync_dir=SYNC_DIR):
        self.base_sync_dir = base_sync_dir
        if not os.path.exists(self.base_sync_dir):
            os.makedirs(self.base_sync_dir, exist_ok=True)

    def get_session_dir(self, session_id: str) -> str:
        """Returns the unique identifier directory for this session's sync."""
        path = os.path.join(self.base_sync_dir, session_id)
        os.makedirs(path, exist_ok=True)
        return path

    def handle_manifest(self, session_id: str, manifest: agent_pb2.DirectoryManifest):
        """Prepares the local directory based on the server manifest."""
        session_dir = self.get_session_dir(session_id)
        print(f"[📁] Preparing Sync Directory: {session_dir}")
        # In Phase 1, we just ensure directories exist
        for file_info in manifest.files:
            if file_info.is_dir:
                os.makedirs(os.path.join(session_dir, file_info.path), exist_ok=True)

    def write_chunk(self, session_id: str, payload: agent_pb2.FilePayload) -> bool:
        """Writes a file chunk to the local session directory."""
        session_dir = self.get_session_dir(session_id)
        target_path = os.path.normpath(os.path.join(session_dir, payload.path))
        
        if not target_path.startswith(session_dir):
            return False # Path traversal guard
            
        os.makedirs(os.path.dirname(target_path), exist_ok=True)
        
        mode = "ab" if payload.chunk_index > 0 else "wb"
        with open(target_path, mode) as f:
            f.write(payload.chunk)
            
        if payload.is_final and payload.hash:
            return self._verify(target_path, payload.hash)
        return True

    def _verify(self, path, expected_hash):
        with open(path, "rb") as f:
            actual = hashlib.sha256(f.read()).hexdigest()
        if actual != expected_hash:
            print(f"[⚠️] Sync Hash Mismatch for {path}")
            return False
        return True