import digitalocean
import os
import time
from dotenv import load_dotenv
from pathlib import Path
from fabric import Connection

# Internal imports
from .utils import get_project_context
from .recipes import generate_stack
from .dockerizer import detect_framework, generate_deployment_assets


class InfraEngine:
    """
    The Core SDK.
    This class handles all communication with DigitalOcean.
    It is designed to be used by a CLI or an AI Agent.
    """

    def __init__(self):
        # Load secrets immediately
        load_dotenv()
        self.token = os.getenv("DIGITAL_OCEAN_TOKEN")
        
        if not self.token:
            raise ValueError("❌ FATAL: No DIGITAL_OCEAN_TOKEN found in .env or environment.")
            
        self.manager = digitalocean.Manager(token=self.token)

    def get_user_info(self):
        """Returns the current user's email and account status."""
        try:
            account = self.manager.get_account()
            return {"email": account.email, "status": account.status, "limit": account.droplet_limit}
        except Exception as e:
            raise ConnectionError(f"Could not connect to DigitalOcean: {e}")

    def _ensure_ssh_key(self):
        """
        Private method.
        1. Dynamically finds local public SSH keys in ~/.ssh/.
        2. Checks if any of them exist on DigitalOcean.
        3. If not, uploads the first one found.
        Returns: The SSHKey object.
        """
        # 1. Find Local Keys
        ssh_dir = Path.home() / ".ssh"
        pub_keys = list(ssh_dir.glob("*.pub"))

        if not pub_keys:
            raise FileNotFoundError(f"No public SSH keys found in {ssh_dir}. Run 'ssh-keygen' first.")

        # 2. Check Remote Keys against Local Keys
        remote_keys = self.manager.get_all_sshkeys()
        remote_key_map = {k.public_key.strip(): k for k in remote_keys}

        first_local_key_content = None
        for key_path in pub_keys:
            local_pub_key = key_path.read_text().strip()
            if not first_local_key_content:
                first_local_key_content = local_pub_key
            
            if local_pub_key in remote_key_map:
                print(f"✨ Found existing SSH key on DigitalOcean: {key_path.name}")
                return remote_key_map[local_pub_key]

        # 3. If no matches, upload the first key found
        if not first_local_key_content:
             raise FileNotFoundError(f"Could not read any public keys in {ssh_dir}.")

        print(f"✨ No matching SSH key found on DigitalOcean. Uploading new key from {pub_keys[0].name}...")
        new_key = digitalocean.SSHKey(token=self.token,
                                      name=f"xenfra-{pub_keys[0].name}-{int(time.time())}",
                                      public_key=first_local_key_content)
        new_key.create()
        return new_key

    def _get_connection(self, ip_address: str):
        """Creates a standardized Fabric connection."""
        return Connection(
            host=ip_address,
            user="root",
            connect_kwargs={
                "timeout": 10
            }
        )

    def _is_setup_complete(self, ip_address: str) -> bool:
        """Checks for the setup completion marker file on the remote server."""
        try:
            c = self._get_connection(ip_address)
            if c.run("test -f /root/setup_complete", warn=True, hide=True).ok:
                return True
        except Exception:
            # This can happen if the server is not ready for SSH
            return False
        return False

    def _auto_heal_server(self, ip_address: str):
        """
        Diagnoses issues and attempts to fix them remotely.
        This is the 'nuclear option' for a stalled setup.
        """
        print(f"\n[ENGINE] 🚑 INITIATING AUTO-HEAL PROTOCOL for {ip_address}...")
        c = self._get_connection(ip_address)

        # A. STOP SERVICES (Prevent Respawn)
        print("[ENGINE] 🛑 Stopping background update services...")
        c.run("systemctl stop unattended-upgrades.service", warn=True, hide=True)
        c.run("systemctl stop apt-daily.service", warn=True, hide=True)
        c.run("systemctl stop apt-daily-upgrade.service", warn=True, hide=True)

        # B. KILL PROCESSES
        print("[ENGINE] 🔪 Killing stuck processes...")
        c.run("pkill apt", warn=True, hide=True)
        c.run("pkill apt-get", warn=True, hide=True)
        c.run("pkill dpkg", warn=True, hide=True)

        # C. REMOVE LOCKS
        print("[ENGINE] PWN: Removing lock files...")
        c.run("rm -f /var/lib/dpkg/lock*", warn=True, hide=True)
        c.run("rm -f /var/lib/apt/lists/lock", warn=True, hide=True)
        c.run("rm -f /var/cache/apt/archives/lock", warn=True, hide=True)

        # D. REPAIR DPKG
        print("[ENGINE] 🔧 Repairing package database...")
        c.run("dpkg --configure -a", warn=True, hide=True)

        # E. RE-MARK AS COMPLETE (to unblock the health check)
        c.run("touch /root/setup_complete", warn=True, hide=True)
        print("[ENGINE] ✅ Auto-Heal Complete.")

    def upload_code(self, ip_address: str) -> bool:
        """Uses Fabric to upload project files to the /root/app directory."""
        print(f"\n[ENGINE] 🚁 INITIATING CODE AIRLIFT to {ip_address}...")
        
        ignored = {'.git', '.venv', 'venv', '__pycache__', '.pytest_cache', '.DS_Store', '.env'}
        
        try:
            c = self._get_connection(ip_address)
            
            print("[ENGINE] 📤 Uploading Project Files...")
            count = 0
            for root, dirs, files in os.walk("."):
                dirs[:] = [d for d in dirs if d not in ignored]
                
                for file in files:
                    if file.endswith(".pyc"): continue
                    
                    local_path = os.path.join(root, file)
                    rel_path = os.path.relpath(local_path, ".")
                    remote_path = f"/root/app/{rel_path}"
                    
                    remote_dir = os.path.dirname(remote_path)
                    c.run(f"mkdir -p {remote_dir}", hide=True)
                    
                    c.put(local_path, remote_path)
                    count += 1

            print(f"[ENGINE] ✅ Airlift Complete: {count} files transferred.")
            return True

        except Exception as e:
            print(f"[ENGINE] ❌ Upload failed: {e}")
            print("[ENGINE] 🔑 Ensure your SSH Agent is running and keys are added.")
            return False

    def list_servers(self):
        """Returns a list of all active servers."""
        return self.manager.get_all_droplets()

    def deploy_server(self, name, region="blr1", size="s-1vcpu-1gb", image="ubuntu-24-04-x64", logger=None):
        """
        Provisions a new server, runs setup, and deploys the application.
        This is a stateful, blocking method that orchestrates the entire process.
        """
        def log(msg):
            if logger:
                logger(f"[ENGINE] {msg}")
            else:
                print(f"[ENGINE] {msg}")

        log(f"🚀 Starting deployment for '{name}'...")

        try:
            # 1. Project Context & Asset Generation
            log("🔎 Detecting project context...")
            context = get_project_context()
            framework, _, _ = detect_framework()
            is_dockerized = False
            if framework:
                log(f"✨ Web framework '{framework}' detected.")
                generate_deployment_assets(context)
                is_dockerized = True

            # 2. Get SSH Key
            ssh_key = self._ensure_ssh_key()

            # 3. Generate Setup Script
            log("📦 Generating cloud-init setup script...")
            setup_script = generate_stack(context, is_dockerized=is_dockerized)

            # 4. Create Droplet
            log(f"☁️  Provisioning droplet '{name}'...")
            droplet = digitalocean.Droplet(token=self.token,
                                           name=name,
                                           region=region,
                                           image=image,
                                           size_slug=size,
                                           ssh_keys=[ssh_key],
                                           user_data=setup_script)
            droplet.create()

            # 5. Wait for Droplet to become active and get IP
            log("⏳ Waiting for droplet to become active and get IP address...")
            while not droplet.ip_address:
                droplet.load()
                time.sleep(5)
            log(f"🌍 Droplet is active at {droplet.ip_address}")

            # 6. Health Check & Auto-Heal Loop
            log("🩺 Performing health checks on the new server...")
            attempts = 0
            max_retries = 18 # ~90 seconds
            is_healthy = False
            while attempts < max_retries:
                if self._is_setup_complete(droplet.ip_address):
                    log("✅ Server setup complete and healthy.")
                    is_healthy = True
                    break
                log(f"   - Still waiting for setup to complete... (attempt {attempts+1}/{max_retries})")
                time.sleep(5)
                attempts += 1
            
            if not is_healthy:
                log("⚠️ Server setup stalled. Initiating auto-heal protocol...")
                self._auto_heal_server(droplet.ip_address)
                if not self._is_setup_complete(droplet.ip_address):
                     raise Exception("Deployment failed: Auto-heal could not recover the server.")
                log("✅ Server recovered successfully after auto-heal.")

            # 7. Upload Code
            self.upload_code(droplet.ip_address)

            # 8. Run Docker Compose if applicable
            if is_dockerized:
                log("🚀 Starting application with Docker Compose...")
                c = self._get_connection(droplet.ip_address)
                c.run("cd /root/app && docker compose up -d", hide=True)

            log("🎉 Deployment successful!")
            return {"status": "success", "ip": droplet.ip_address, "name": name}

        except Exception as e:
            log(f"❌ CRITICAL DEPLOYMENT FAILURE: {e}")
            # Optionally, trigger a cleanup action here
            raise e

    def destroy_server(self, droplet_id):
        """Destroys a server by ID. IRREVERSIBLE."""
        droplet = self.manager.get_droplet(droplet_id)
        droplet.destroy()
        return True
