"""Multi-GPU sync manager for model distribution."""

from __future__ import annotations

import subprocess
import os
from pathlib import Path
from datetime import datetime
from typing import Optional, List, Dict, Any, Callable, Tuple
from dataclasses import dataclass, field

from .metadata_store import MetadataStore


# VRAM requirement estimates for common model patterns
VRAM_ESTIMATES: Dict[str, int] = {
    # Large video models
    "wan": 48000,
    "hunyuan": 48000,
    "mochi": 48000,
    "cogvideo": 48000,

    # Flux variants
    "flux-dev": 24000,
    "flux_dev": 24000,
    "flux-schnell": 12000,
    "flux_schnell": 12000,
    "flux-fill": 24000,

    # SD variants
    "sdxl": 12000,
    "sd_xl": 12000,
    "sd3": 16000,
    "sd15": 6000,
    "sd1.5": 6000,
    "sd-1": 6000,

    # Other large models
    "cascade": 20000,
    "pixart": 16000,
    "playground": 16000,

    # LoRA typically need less
    "lora": 2000,
}


@dataclass
class SyncResult:
    """Result of a sync operation."""
    target_name: str
    success: bool
    models_synced: int = 0
    models_skipped: int = 0
    bytes_transferred: int = 0
    duration_seconds: float = 0.0
    error: Optional[str] = None
    details: List[str] = field(default_factory=list)


class SyncManager:
    """Manages model synchronization to remote machines."""

    def __init__(
        self,
        store: MetadataStore,
        base_model_path: Optional[str] = None,
        on_progress: Optional[Callable[[str, float], None]] = None,
    ):
        """
        Initialize the sync manager.

        Args:
            store: MetadataStore instance
            base_model_path: Base path for ComfyUI models directory
            on_progress: Callback(target_name, progress) for progress updates
        """
        self.store = store
        self.base_model_path = Path(base_model_path) if base_model_path else None
        self.on_progress = on_progress

    def estimate_vram(self, model: Dict[str, Any]) -> int:
        """
        Estimate VRAM requirement for a model.

        Args:
            model: Model dict from database

        Returns:
            Estimated VRAM in MB
        """
        filename = model.get("filename", "").lower()
        model_type = model.get("model_type", "").lower()

        # Check filename patterns
        for pattern, vram in VRAM_ESTIMATES.items():
            if pattern in filename:
                return vram

        # Check if it's a LoRA by type
        if "lora" in model_type:
            return VRAM_ESTIMATES.get("lora", 2000)

        # Size-based fallback
        file_size = model.get("file_size", 0)
        if file_size > 0:
            size_gb = file_size / (1024 ** 3)
            # Rough estimate: 2GB VRAM per GB of model file
            return int(size_gb * 2000)

        # Default: assume it needs moderate VRAM
        return 8000

    def get_models_for_target(
        self,
        target: Dict[str, Any],
        models: Optional[List[Dict[str, Any]]] = None
    ) -> List[Dict[str, Any]]:
        """
        Get models filtered by target's VRAM limit.

        Args:
            target: Sync target dict
            models: Optional list of models (fetches all if not provided)

        Returns:
            List of models that fit target's VRAM limit
        """
        if models is None:
            models = self.store.get_all_models()

        max_vram = target.get("max_vram_mb")

        filtered = []
        for model in models:
            # Skip if explicitly excluded
            if model.get("exclude_from_sync"):
                continue

            # Check VRAM requirement
            if max_vram is not None:
                model_vram = model.get("vram_mb") or self.estimate_vram(model)
                if model_vram > max_vram:
                    continue

            filtered.append(model)

        return filtered

    def get_relative_path(self, model: Dict[str, Any]) -> str:
        """
        Get the relative path for a model within the models directory.

        Args:
            model: Model dict

        Returns:
            Relative path string
        """
        model_path = Path(model["path"])

        if self.base_model_path:
            try:
                return str(model_path.relative_to(self.base_model_path))
            except ValueError:
                pass

        # Fallback: try to extract from common patterns
        # e.g., /path/to/ComfyUI/models/checkpoints/model.safetensors
        parts = model_path.parts
        for i, part in enumerate(parts):
            if part == "models" and i < len(parts) - 1:
                return str(Path(*parts[i + 1:]))

        # Last resort: just the filename
        return model_path.name

    def sync_to_target(
        self,
        target: Dict[str, Any],
        models: Optional[List[Dict[str, Any]]] = None,
        dry_run: bool = False
    ) -> SyncResult:
        """
        Sync models to a target machine using rsync.

        Args:
            target: Sync target dict
            models: Optional list of models (uses filtered list if not provided)
            dry_run: If True, only simulate the sync

        Returns:
            SyncResult with details
        """
        start_time = datetime.now()
        result = SyncResult(target_name=target["name"], success=True)

        if models is None:
            models = self.get_models_for_target(target)

        if not models:
            result.details.append("No models to sync")
            return result

        ssh_host = target.get("ssh_host")
        target_path = target["model_path"]

        for model in models:
            rel_path = self.get_relative_path(model)
            model_path = model["path"]

            # Build destination path
            if ssh_host:
                dest = f"{ssh_host}:{target_path}/{rel_path}"
            else:
                dest = f"{target_path}/{rel_path}"

            # Ensure destination directory exists
            dest_dir = Path(rel_path).parent
            if ssh_host:
                mkdir_cmd = ["ssh", ssh_host, f"mkdir -p '{target_path}/{dest_dir}'"]
            else:
                mkdir_cmd = ["mkdir", "-p", f"{target_path}/{dest_dir}"]

            try:
                if not dry_run:
                    subprocess.run(mkdir_cmd, check=True, capture_output=True)
            except subprocess.CalledProcessError as e:
                result.details.append(f"Failed to create directory for {rel_path}: {e}")
                result.models_skipped += 1
                continue

            # Build rsync command
            rsync_args = [
                "rsync",
                "-avz",  # archive, verbose, compress
                "--progress",
            ]

            if dry_run:
                rsync_args.append("--dry-run")

            # Add SSH options if remote
            if ssh_host:
                rsync_args.extend(["-e", "ssh"])

            rsync_args.extend([model_path, dest])

            try:
                proc = subprocess.run(
                    rsync_args,
                    capture_output=True,
                    text=True,
                    timeout=3600  # 1 hour timeout per file
                )

                if proc.returncode == 0:
                    result.models_synced += 1
                    # Parse transferred bytes from output
                    # rsync output: "sent X bytes  received Y bytes"
                    for line in proc.stdout.split("\n"):
                        if "sent" in line and "bytes" in line:
                            try:
                                sent = int(line.split("sent")[1].split("bytes")[0].strip().replace(",", ""))
                                result.bytes_transferred += sent
                            except (ValueError, IndexError):
                                pass
                else:
                    result.details.append(f"rsync failed for {rel_path}: {proc.stderr}")
                    result.models_skipped += 1

            except subprocess.TimeoutExpired:
                result.details.append(f"Timeout syncing {rel_path}")
                result.models_skipped += 1
            except Exception as e:
                result.details.append(f"Error syncing {rel_path}: {e}")
                result.models_skipped += 1

            # Report progress
            if self.on_progress:
                progress = (result.models_synced + result.models_skipped) / len(models)
                self.on_progress(target["name"], progress)

        # Update last sync time
        if not dry_run and result.models_synced > 0:
            self.store.update_sync_target(target["id"], {
                "last_sync": datetime.now().isoformat()
            })

        result.duration_seconds = (datetime.now() - start_time).total_seconds()
        result.success = result.models_skipped == 0

        return result

    def sync_to_all_targets(
        self,
        dry_run: bool = False
    ) -> List[SyncResult]:
        """
        Sync models to all enabled targets.

        Args:
            dry_run: If True, only simulate the sync

        Returns:
            List of SyncResults for each target
        """
        targets = self.store.get_all_sync_targets(enabled_only=True)
        results = []

        for target in targets:
            result = self.sync_to_target(target, dry_run=dry_run)
            results.append(result)

        return results

    def generate_sync_script(
        self,
        target: Dict[str, Any],
        models: Optional[List[Dict[str, Any]]] = None
    ) -> str:
        """
        Generate a standalone bash script for syncing.

        Args:
            target: Sync target dict
            models: Optional list of models

        Returns:
            Bash script content
        """
        if models is None:
            models = self.get_models_for_target(target)

        ssh_host = target.get("ssh_host")
        target_path = target["model_path"]

        lines = [
            "#!/bin/bash",
            "# ModelSync script for target: " + target["name"],
            f"# Generated: {datetime.now().isoformat()}",
            f"# Models: {len(models)}",
            "",
            "set -e",
            "",
            f'TARGET_PATH="{target_path}"',
        ]

        if ssh_host:
            lines.append(f'SSH_HOST="{ssh_host}"')
            lines.append("")
            lines.append("# Ensure remote directories exist")
        else:
            lines.append("")
            lines.append("# Ensure directories exist")

        # Collect unique directories
        dirs = set()
        for model in models:
            rel_path = self.get_relative_path(model)
            parent = str(Path(rel_path).parent)
            if parent and parent != ".":
                dirs.add(parent)

        for d in sorted(dirs):
            if ssh_host:
                lines.append(f'ssh "$SSH_HOST" "mkdir -p \'$TARGET_PATH/{d}\'"')
            else:
                lines.append(f'mkdir -p "$TARGET_PATH/{d}"')

        lines.append("")
        lines.append("# Sync models")

        for model in models:
            rel_path = self.get_relative_path(model)
            src = model["path"]

            if ssh_host:
                lines.append(
                    f'rsync -avz --progress "{src}" "$SSH_HOST:$TARGET_PATH/{rel_path}"'
                )
            else:
                lines.append(
                    f'rsync -avz --progress "{src}" "$TARGET_PATH/{rel_path}"'
                )

        lines.append("")
        lines.append("echo 'Sync complete!'")

        return "\n".join(lines)

    def verify_target_connectivity(self, target: Dict[str, Any]) -> Tuple[bool, str]:
        """
        Verify SSH connectivity to a target.

        Args:
            target: Sync target dict

        Returns:
            Tuple of (success, message)
        """
        ssh_host = target.get("ssh_host")

        if not ssh_host:
            # Local target, check if path exists
            target_path = Path(target["model_path"])
            if target_path.exists():
                return True, f"Local path exists: {target_path}"
            else:
                return False, f"Local path does not exist: {target_path}"

        try:
            result = subprocess.run(
                ["ssh", "-o", "ConnectTimeout=10", ssh_host, "echo", "OK"],
                capture_output=True,
                text=True,
                timeout=15
            )

            if result.returncode == 0:
                return True, f"SSH connection successful to {ssh_host}"
            else:
                return False, f"SSH connection failed: {result.stderr}"

        except subprocess.TimeoutExpired:
            return False, f"SSH connection timeout to {ssh_host}"
        except Exception as e:
            return False, f"Error connecting to {ssh_host}: {e}"

    def get_sync_summary(self) -> Dict[str, Any]:
        """
        Get a summary of sync status across all targets.

        Returns:
            Dict with summary information
        """
        targets = self.store.get_all_sync_targets()
        all_models = self.store.get_all_models()

        summary = {
            "total_models": len(all_models),
            "targets": [],
        }

        for target in targets:
            eligible_models = self.get_models_for_target(target, all_models)

            summary["targets"].append({
                "name": target["name"],
                "enabled": bool(target.get("enabled")),
                "max_vram_mb": target.get("max_vram_mb"),
                "eligible_models": len(eligible_models),
                "last_sync": target.get("last_sync"),
            })

        return summary