import os
import shutil
import hashlib
import time
import json
import zlib
from typing import Dict, List
from app.core.grpc.shared_core.ignore import CortexIgnore
from mesh_core import agent_pb2
from app.config import settings

class GhostMirrorManager:
    """Manages local server-side copies of node workspaces."""
    def __init__(self, storage_root="/app/data/mirrors"):
        self.storage_root = os.path.abspath(storage_root)
        if not os.path.exists(self.storage_root):
            os.makedirs(self.storage_root, exist_ok=True)
        self.hash_cache = {} # (abs_path) -> (size, mtime, hash)
            
    def get_ignore_filter(self, session_id: str) -> CortexIgnore:
        """Returns a CortexIgnore instance for a session."""
        return CortexIgnore(self.get_workspace_path(session_id))

    def get_workspace_path(self, session_id: str, create: bool = True) -> str:
        """Returns the local absolute path for a session's mirror."""
        path = os.path.join(self.storage_root, session_id)
        if create:
            os.makedirs(path, exist_ok=True)
        return path

    def purge(self, session_id: str):
        """Completely removes a session's mirror from disk."""
        path = os.path.join(self.storage_root, session_id)
        if os.path.exists(path):
            shutil.rmtree(path)
            print(f"    [📁🧹] Ghost Mirror deleted for session: {session_id}")

    def write_file_chunk(self, session_id: str, file_payload: agent_pb2.FilePayload):
        """Writes a chunk of data to the local mirror."""
        workspace = self.get_workspace_path(session_id)
        
        # Phase 3 ignore filter
        ignore_filter = CortexIgnore(workspace, is_upstream=True)
        if ignore_filter.is_ignored(file_payload.path):
            print(f"    [📁🚷] Ignoring write to {file_payload.path}")
            return

        if file_payload.path.endswith(".cortex_tmp") or file_payload.path.endswith(".cortex_lock"):
            return

        # Prevent path traversal — strip leading slash to ensure it's relative to workspace
        path_safe = file_payload.path.lstrip("/")
        abs_workspace = os.path.abspath(workspace)
        safe_path = os.path.abspath(os.path.join(abs_workspace, path_safe))
        if not safe_path.startswith(abs_workspace):
            raise ValueError(f"Malicious path detected: file={file_payload.path} safe_path={safe_path} workspace={abs_workspace}")

        # --- SYSTEM SKILL IMMUTABILITY LOCK ---
        # Prevent any AI agent from modifying strict system skills across the mesh
        if path_safe.startswith(".skills/"):
            parts = path_safe.split("/")
            if len(parts) >= 2:
                skill_dir = os.path.join(workspace, ".skills", parts[1])
                if os.path.islink(skill_dir):
                    try:
                        meta_path = os.path.join(os.path.realpath(skill_dir), ".metadata.json")
                        if os.path.exists(meta_path):
                            with open(meta_path, "r") as mf:
                                meta_json = json.load(mf)
                                if meta_json.get("is_system", False):
                                    print(f"    [📁🚫] SYSTEM IMMUTABILITY LOCK: Node blocked from modifying immutable system skill '{parts[1]}'.")
                                    return
                    except Exception:
                        pass
        # ----------------------------------------
                        
        os.makedirs(os.path.dirname(safe_path), exist_ok=True)
        
        # 0. Fast-Path for single-chunk files (Zero UX latency, no locking)
        total_chunks = getattr(file_payload, "total_chunks", 0)
        if total_chunks == 1 and file_payload.chunk_index == 0 and file_payload.is_final:
            data = file_payload.chunk
            if file_payload.compressed:
                try: data = zlib.decompress(data)
                except: pass
            
            with open(safe_path, "wb") as f:
                f.write(data)
                
            success = True
            if file_payload.hash:
                success = self._verify_hash(safe_path, file_payload.hash)
            
            if not success:
                try: os.remove(safe_path)
                except: pass
                print(f"    [📁❌] Fast Sync failed hash check for {file_payload.path}")
            else:
                try:
                    parent_stat = os.stat(os.path.dirname(safe_path))
                    os.chown(safe_path, parent_stat.st_uid, parent_stat.st_gid)
                except: pass
                print(f"    [📁⚡] Fast Sync Complete: {file_payload.path}")
                self._check_skill_promotion(session_id, file_payload.path, safe_path)
            return
            
        tmp_path = safe_path + ".cortex_tmp"
        lock_path = safe_path + ".cortex_lock"
        
        # 1. Lock/Metadata initialization or update
        if file_payload.chunk_index == 0:
            print(f"    [📁📥] Sync Starting: {file_payload.path} ({file_payload.total_size} bytes, {file_payload.total_chunks} chunks)")
            try:
                metadata = {
                    "ts": time.time(),
                    "owner": "hub",
                    "path": file_payload.path,
                    "total_chunks": getattr(file_payload, "total_chunks", 0),
                    "total_size": getattr(file_payload, "total_size", 0),
                    "received_chunks": [0]
                }
                with open(lock_path, "w") as lf:
                    lf.write(json.dumps(metadata))
            except: pass
        elif file_payload.chunk_index % 10 == 0:
            try:
                if os.path.exists(lock_path):
                    with open(lock_path, "r+") as lf:
                        meta = json.loads(lf.read())
                        if "received_chunks" not in meta: meta["received_chunks"] = []
                        if file_payload.chunk_index not in meta["received_chunks"]:
                            meta["received_chunks"].append(file_payload.chunk_index)
                            lf.seek(0)
                            lf.write(json.dumps(meta))
                            lf.truncate()
            except: pass

        # 2. Decompress and Write Data (handles out-of-order chunks)
        data = file_payload.chunk
        if file_payload.compressed:
            try: data = zlib.decompress(data)
            except: pass

        # Open in r+b if exists to avoid truncating chunks that might have arrived out-of-order,
        # otherwise create with wb.
        mode = "r+b" if os.path.exists(tmp_path) else "wb"
        with open(tmp_path, mode) as f:
            f.seek(file_payload.offset)
            f.write(data)
            
        if file_payload.is_final:
            success = True
            if file_payload.hash:
                success = self._verify_hash(tmp_path, file_payload.hash)
            
            if success:
                try:
                    # M6: Ownership Preservation
                    try:
                        parent_stat = os.stat(os.path.dirname(safe_path))
                        os.chown(tmp_path, parent_stat.st_uid, parent_stat.st_gid)
                    except: pass

                    os.replace(tmp_path, safe_path)
                    print(f"    [📁✅] Sync Complete: {file_payload.path} (Swapped and verified)")
                    self._check_skill_promotion(session_id, file_payload.path, safe_path)
                except Exception as e:
                    print(f"    [📁❌] Server atomic swap failed for {file_payload.path}: {e}")
                    success = False
            
            if os.path.exists(lock_path):
                try: os.remove(lock_path)
                except: pass
            if os.path.exists(tmp_path) and not success:
                try: os.remove(tmp_path)
                except: pass

    def _check_skill_promotion(self, session_id: str, rel_path: str, safe_path: str):
        import re
        import shutil
        match = re.match(r"^\.skills/([^/]+)/[sS][kK][iI][lL][lL]\.md$", rel_path)
        if match:
            skill_name = match.group(1)
            
            # Normalize filename explicitly to uppercase SKILL.md for the UI backend loader
            if not safe_path.endswith("SKILL.md"):
                new_safe_path = os.path.join(os.path.dirname(safe_path), "SKILL.md")
                try: os.rename(safe_path, new_safe_path)
                except: pass
                safe_path = new_safe_path
            skill_workspace_dir = os.path.dirname(safe_path)
            
            # If it's a real folder and not a symlink, the AI created it locally
            if os.path.isdir(skill_workspace_dir) and not os.path.islink(skill_workspace_dir):
                from app.config import settings
                feature_name = "swarm_control"
                owner_id = "admin"
                
                # Attempt to determine feature and owner from session DB
                try:
                    from app.db.session import get_db_session
                    from app.db.models import Session
                    with get_db_session() as db:
                        sess = db.query(Session).filter(Session.sync_workspace_id == session_id).first()
                        if sess:
                            owner_id = sess.user_id
                            if sess.feature_name:
                                feature_name = sess.feature_name
                except Exception:
                    pass
                
                global_skills_dir = os.path.join(settings.DATA_DIR, "skills", feature_name)
                os.makedirs(global_skills_dir, exist_ok=True)
                target_physical_dir = os.path.join(global_skills_dir, skill_name)
                
                if not os.path.exists(target_physical_dir):
                    try:
                        print(f"    [📁⭐] Symlink Inversion Promotor Triggered for: {skill_name}")
                        # Move physical files securely into backend VFS
                        shutil.move(skill_workspace_dir, target_physical_dir)
                        
                        # Apply Metadata
                        meta_path = os.path.join(target_physical_dir, ".metadata.json")
                        if not os.path.exists(meta_path):
                            meta = {
                                "owner_id": owner_id,
                                "is_system": False,
                                "extra_metadata": {"emoji": "🤖"}
                            }
                            with open(meta_path, "w") as f:
                                json.dump(meta, f, indent=4)
                        
                        # Plop a symlink back into the session sandbox
                        os.symlink(target_physical_dir, skill_workspace_dir, target_is_directory=True)
                        print(f"    [📁⭐] Successfully promoted {skill_name} to global VFS!")
                    except Exception as e:
                        print(f"    [📁❌] Failed to promote local skill {skill_name}: {e}")

    def delete_file(self, session_id: str, rel_path: str):
        """Deletes a file or directory from the local mirror."""
        workspace = self.get_workspace_path(session_id, create=False)
        if not os.path.exists(workspace):
            return  # Parent folder is already gone

        path_safe = rel_path.lstrip("/")
        abs_workspace = os.path.abspath(workspace)
        safe_path = os.path.abspath(os.path.join(abs_workspace, path_safe))
        
        if not safe_path.startswith(abs_workspace):
            raise ValueError(f"Malicious path detected: {rel_path}")

        # --- SYSTEM SKILL IMMUTABILITY LOCK ---
        if path_safe.startswith(".skills/"):
            parts = path_safe.split("/")
            if len(parts) >= 2:
                skill_dir = os.path.join(workspace, ".skills", parts[1])
                if os.path.islink(skill_dir):
                    try:
                        meta_path = os.path.join(os.path.realpath(skill_dir), ".metadata.json")
                        if os.path.exists(meta_path):
                            with open(meta_path, "r") as mf:
                                meta_json = json.load(mf)
                                if meta_json.get("is_system", False):
                                    print(f"    [📁🚫] SYSTEM IMMUTABILITY LOCK: Node blocked from deleting immutable system skill '{parts[1]}'.")
                                    return
                    except Exception:
                        pass
        # ----------------------------------------
            
        if os.path.exists(safe_path) or os.path.islink(safe_path):
            if os.path.islink(safe_path):
                os.unlink(safe_path)
            elif os.path.isdir(safe_path):
                shutil.rmtree(safe_path)
            else:
                os.remove(safe_path)
            print(f"    [📁🗑️] Mirror Sync: Deleted {rel_path}")
    
    def move_file(self, session_id: str, old_rel_path: str, new_rel_path: str):
        """Atomically moves/renames a file or directory in the mirror."""
        workspace = self.get_workspace_path(session_id)
        
        old_path = os.path.abspath(os.path.join(workspace, old_rel_path.lstrip("/")))
        new_path = os.path.abspath(os.path.join(workspace, new_rel_path.lstrip("/")))
        
        # Security: Prevent path traversal
        abs_workspace = os.path.abspath(workspace)
        if not old_path.startswith(abs_workspace) or not new_path.startswith(abs_workspace):
             raise ValueError("Path traversal detected in move")
             
        # Security: Immutability lock for skills
        if old_rel_path.startswith(".skills/") or new_rel_path.startswith(".skills/"):
            # Simple check for now: don't move things into or out of .skills/ if they are system
             pass # Logic from delete_file/write_file can be applied if needed
             
        if os.path.exists(old_path):
            os.makedirs(os.path.dirname(new_path), exist_ok=True)
            shutil.move(old_path, new_path)
            print(f"    [📁🚚] Mirror Sync: Moved {old_rel_path} -> {new_rel_path}")

    def copy_file(self, session_id: str, old_rel_path: str, new_rel_path: str):
        """Copies a file or directory in the mirror."""
        workspace = self.get_workspace_path(session_id)
        
        old_path = os.path.abspath(os.path.join(workspace, old_rel_path.lstrip("/")))
        new_path = os.path.abspath(os.path.join(workspace, new_rel_path.lstrip("/")))
        
        abs_workspace = os.path.abspath(workspace)
        if not old_path.startswith(abs_workspace) or not new_path.startswith(abs_workspace):
            raise ValueError("Path traversal detected in copy")

        if os.path.exists(old_path):
            os.makedirs(os.path.dirname(new_path), exist_ok=True)
            if os.path.isdir(old_path):
                shutil.copytree(old_path, new_path, dirs_exist_ok=True)
            else:
                shutil.copy2(old_path, new_path)
            print(f"    [📁👯] Mirror Sync: Copied {old_rel_path} -> {new_rel_path}")

    def stat_file(self, session_id: str, rel_path: str) -> dict:
        """Returns metadata for a file in the Hub mirror."""
        workspace = self.get_workspace_path(session_id)
        path = os.path.abspath(os.path.join(workspace, rel_path.lstrip("/")))
        
        abs_workspace = os.path.abspath(workspace)
        if not path.startswith(abs_workspace):
             return {"error": f"Path traversal detected: {path} vs {abs_workspace}"}

        if not os.path.exists(path):
            return {"error": "Not Found"}

        s = os.stat(path)
        return {
            "path": rel_path,
            "name": os.path.basename(path),
            "exists": True,
            "size": s.st_size,
            "mtime": s.st_mtime,
            "is_dir": os.path.isdir(path),
            "is_file": os.path.isfile(path),
            "is_link": os.path.islink(path)
        }

    def generate_manifest(self, session_id: str) -> agent_pb2.DirectoryManifest:
        """Generates a manifest of the current local mirror state using parallel hashing."""
        from concurrent.futures import ThreadPoolExecutor
        workspace = self.get_workspace_path(session_id)
        ignore_filter = self.get_ignore_filter(session_id)

        raw_files = []
        for root, dirs, filenames in os.walk(workspace, followlinks=True):
            dirs[:] = [d for d in dirs if not ignore_filter.is_ignored(os.path.relpath(os.path.join(root, d), workspace))]
            for filename in filenames:
                abs_path = os.path.join(root, filename)
                rel_path = os.path.relpath(abs_path, workspace)
                if ignore_filter.is_ignored(rel_path):
                    continue
                raw_files.append((abs_path, rel_path))

        def _hash_file(file_tuple):
            abs_path, rel_path = file_tuple
            try:
                stats = os.stat(abs_path)
                mtime = stats.st_mtime
                size = stats.st_size

                cached = self.hash_cache.get(abs_path)
                if cached and cached[0] == size and cached[1] == mtime:
                    file_hash = cached[2]
                else:
                    h = hashlib.sha256()
                    with open(abs_path, "rb") as f:
                        while True:
                            chunk = f.read(1024 * 1024)  # 1MB chunks
                            if not chunk:
                                break
                            h.update(chunk)
                    file_hash = h.hexdigest()
                    self.hash_cache[abs_path] = (size, mtime, file_hash)

                return agent_pb2.FileInfo(
                    path=rel_path,
                    size=size,
                    hash=file_hash,
                    is_dir=False,
                )
            except Exception:
                return None

        files = []
        with ThreadPoolExecutor(max_workers=min(settings.MIRROR_MAX_WORKERS, len(raw_files) or 1)) as executor:
            results = list(executor.map(_hash_file, raw_files))
            files = [r for r in results if r is not None]

        return agent_pb2.DirectoryManifest(root_path=workspace, files=files)

    def _verify_hash(self, path: str, expected_hash: str) -> bool:
        if not os.path.exists(path): return False
        try:
            h = hashlib.sha256()
            with open(path, "rb") as f:
                while True:
                    chunk = f.read(1024 * 1024)
                    if not chunk: break
                    h.update(chunk)
            actual_hash = h.hexdigest()
            if actual_hash != expected_hash:
                print(f"[⚠️] Hash Mismatch for {path}: expected {expected_hash}, got {actual_hash}")
                return False
            return True
        except Exception as e:
            print(f"[⚠️] Hash verification error: {e}")
            return False

    def reconcile(self, session_id: str, remote_manifest: agent_pb2.DirectoryManifest) -> List[str]:
        """Compares remote manifest with local mirror using parallel verification."""
        from concurrent.futures import ThreadPoolExecutor
        workspace = self.get_workspace_path(session_id)
        ignore_filter = CortexIgnore(workspace, is_upstream=True)
        
        expected_paths = {f.path for f in remote_manifest.files}

        # 1. Purge extraneous local files (Legacy remains synchronous as it's metadata-only)
        manifest_root = remote_manifest.root_path or "."
        for root, dirs, files_in_dir in os.walk(workspace, topdown=False):
            for name in files_in_dir:
                abs_path = os.path.join(root, name)
                rel_path = os.path.relpath(abs_path, workspace)
                is_within_root = (manifest_root == "." or rel_path.startswith(manifest_root + os.sep) or rel_path == manifest_root)
                if not is_within_root: continue
                if rel_path in [".cortexignore", ".gitignore"]: continue
                if rel_path not in expected_paths and not ignore_filter.is_ignored(rel_path):
                    try: os.remove(abs_path)
                    except Exception: pass
            
            for name in dirs:
                abs_path = os.path.join(root, name)
                rel_path = os.path.relpath(abs_path, workspace)
                is_within_root = (manifest_root == "." or rel_path.startswith(manifest_root + os.sep) or rel_path == manifest_root)
                if not is_within_root: continue
                if rel_path not in expected_paths and not ignore_filter.is_ignored(rel_path):
                    try: 
                        if not os.listdir(abs_path): os.rmdir(abs_path)
                    except Exception: pass

        # 2. Reconcile with parallel hash checks
        def _check_drift(remote_file):
            if remote_file.is_dir: return None
            local_path = os.path.join(workspace, remote_file.path)
            
            if not os.path.exists(local_path):
                return remote_file.path
            
            # Metadata pre-check (Fast path)
            try:
                stats = os.stat(local_path)
                size = stats.st_size
                mtime = stats.st_mtime
                if size != remote_file.size:
                    return remote_file.path
                
                cached = self.hash_cache.get(local_path)
                if cached and cached[0] == size and cached[1] == mtime:
                    if cached[2] != remote_file.hash:
                        return remote_file.path
                    return None
            except: return remote_file.path

            # Hash check (Slow path) - Incremental
            try:
                h = hashlib.sha256()
                with open(local_path, "rb") as f:
                    while True:
                        chunk = f.read(1024 * 1024)
                        if not chunk: break
                        h.update(chunk)
                file_hash = h.hexdigest()
                self.hash_cache[local_path] = (size, mtime, file_hash)
                if file_hash != remote_file.hash:
                    return remote_file.path
            except: 
                return remote_file.path
            return None

        needs_update = []
        with ThreadPoolExecutor(max_workers=settings.MIRROR_MAX_WORKERS) as executor:
            drift_results = list(executor.map(_check_drift, remote_manifest.files))
            needs_update = [d for d in drift_results if d is not None]
            
        return needs_update

    def purge_orphaned(self, active_session_ids: List[str]):
        """
        Deletes mirror directories that are not in the provided active list.
        Ignores special directories like '__fs_explorer__' or internal bookmarks.
        """
        if not os.path.exists(self.storage_root):
            return

        print(f"    [📁🧹] Running Mirror Cleanup. Active Sessions: {len(active_session_ids)}")
        active_set = set(active_session_ids)
        # Always preserve special system folders
        active_set.add("__fs_explorer__")
        active_set.add(".cortexignore")

        for entry in os.scandir(self.storage_root):
            if entry.is_dir() and entry.name not in active_set:
                try:
                    # Robust recursive delete with small retry for race conditions
                    for i in range(3):
                        try:
                            shutil.rmtree(entry.path)
                            print(f"    [📁🗑️] Purged orphaned ghost mirror: {entry.name}")
                            break
                        except OSError as se:
                            if i == 2: raise se
                            time.sleep(0.5) 
                except Exception as e:
                    print(f"    [📁⚠️] Failed to purge orphaned mirror {entry.name}: {e}")

    def purge_stale_locks(self, session_id: str, ttl: int = 60):
        """Scans a mirror for stale .cortex_lock files and removes them."""
        workspace = self.get_workspace_path(session_id)
        now = time.time()
        purged = 0
        
        for root, _, files in os.walk(workspace):
            for f in files:
                if f.endswith(".cortex_lock"):
                    abs_path = os.path.join(root, f)
                    tmp_path = abs_path.replace(".cortex_lock", ".cortex_tmp")
                    try:
                        with open(abs_path, "r") as lf:
                            data = json.loads(lf.read())
                            if now - data.get("ts", 0) > ttl:
                                if os.path.exists(abs_path): os.remove(abs_path)
                                if os.path.exists(tmp_path): os.remove(tmp_path)
                                purged += 1
                                print(f"    [📁🔓] Purged stale lock/tmp from {data.get('owner')}: {data.get('path')}")
                    except:
                        # If unreadable/corrupt, purge it anyway
                        try: 
                            if os.path.exists(abs_path): os.remove(abs_path)
                            if os.path.exists(tmp_path): os.remove(tmp_path)
                            purged += 1
                        except: pass
                elif f.endswith(".cortex_tmp"):
                    # Orphaned tmp file check
                    abs_path = os.path.join(root, f)
                    lock_path = abs_path.replace(".cortex_tmp", ".cortex_lock")
                    if not os.path.exists(lock_path):
                        try:
                            if now - os.path.getmtime(abs_path) > ttl:
                                os.remove(abs_path)
                                purged += 1
                        except: pass
        if purged > 0:
            print(f"    [📁🧹] Cleaned up {purged} stale locks/tmp files in {session_id}")
