from __future__ import annotations
import logging
import subprocess
import sys
import time
from pathlib import Path
from typing import TYPE_CHECKING

from installer.milestone import Milestone
from installer import config as inst_config

if TYPE_CHECKING:
    from installer.state import InstallerState
    from installer.milestone import ProgressCallback


def _stream_line(line: str) -> None:
    try:
        sys.stdout.write(line)
        sys.stdout.flush()
    except Exception:
        pass

class DependenciesMilestone(Milestone):
    def __init__(self):
        super().__init__("dependencies", "Dependency Verification")
        self.logger = logging.getLogger("DependenciesMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if progress:
            progress("Checking dependencies...", 0.1)
        
        dependencies = inst_config.DEPENDENCIES
        missing = []
        
        for i, dep in enumerate(dependencies):
            ok, location, version = inst_config.get_dep_info(dep)
            if ok:
                self.logger.info(f"[OK] {dep['name']} {version or ''}".strip())
            else:
                missing.append(dep)
                self.logger.info(f"[MISSING] {dep['name']}")
            
            if progress:
                progress(f"Checking {dep['name']}...", 0.1 + (i / len(dependencies)) * 0.4)

        if not missing:
            self._set_status(state, "All installed")
            if progress:
                progress("All dependencies installed", 1.0)
            return

        # Check if auto-install is enabled
        auto_install = self._parse_bool(state.inputs.get('dependencies.auto_install_missing', 'True'))
        if not auto_install:
            self.logger.error("Dependencies missing and auto-install disabled.")
            self._set_status(state, "Missing")
            return

        for i, dep in enumerate(missing):
            install_cmd = dep.get('install_cmd')
            if not install_cmd:
                self.logger.error(f"No install command for {dep['name']}.")
                continue
            
            should_install = self._parse_bool(state.inputs.get(f"dependencies.{dep['id']}.install", 'True'))
            if not should_install:
                self.logger.info(f"[SKIP] {dep['name']} install disabled by config.")
                continue
            
            self.logger.info(f"[INSTALL] {dep['name']} -> {install_cmd}")
            if progress:
                progress(f"Installing {dep['name']}...", 0.5 + (i / len(missing)) * 0.4)
                
            rc = self._run_cmd(install_cmd)
            if rc != 0:
                self.logger.error(f"Install failed for {dep['name']} (code {rc})")

        # Re-verify
        still_missing = []
        for dep in missing:
            ok, _, _ = inst_config.get_dep_info(dep)
            if not ok:
                still_missing.append(dep['name'])
        
        if still_missing:
            self.logger.error(f"Still missing: {', '.join(still_missing)}")
            self._set_status(state, "Missing")
        else:
            self._set_status(state, "All installed")
            if progress:
                progress("All dependencies installed", 1.0)

    def _set_status(self, state: InstallerState, status: str):
        if 'Dependencies' not in state.config_data:
            state.config_data['Dependencies'] = {}
        state.config_data['Dependencies']['STATUS'] = status


class NetworkScanMilestone(Milestone):
    def __init__(self):
        super().__init__("network_scan", "Network Configuration")
        self.logger = logging.getLogger("NetworkScanMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if not self._parse_bool(state.inputs.get('network_scan.run', 'True')):
            self.logger.info("Network scan disabled.")
            return

        existing_kdc = state.config_data.get('Network', {}).get('KDC_AUTO_DETECTED')
        if existing_kdc:
            self.logger.info(f"Network scan already performed. KDC: {existing_kdc}")
            return

        scan_binary = inst_config.get_resource_path("prole-net/prole-agent")
        if not scan_binary.exists():
            self.logger.error(f"Scan binary not found at {scan_binary}")
            return

        prole_home = Path.home() / ".prole"
        scan_dir = prole_home / "scan"
        scan_dir.mkdir(parents=True, exist_ok=True)

        kdc_found = None

        def _handle_stdout(line: str):
            nonlocal kdc_found
            if "KDC is:" in line:
                try:
                    ip_part = line.split("KDC is:")[1].strip()
                    ip = ip_part.split()[0].strip('[]():,')
                    if ip:
                        kdc_found = ip
                except Exception:
                    pass
            elif "Active Directory" in line or "88" in line:
                import socket
                for part in line.split():
                    try:
                        socket.inet_aton(part.strip('[]():,'))
                        kdc_found = part.strip('[]():,')
                        break
                    except Exception:
                        continue

        if progress:
            progress("Scanning network...", 0.5)
            
        rc = self._run_cmd([str(scan_binary), "-t", "10"], cwd=str(scan_dir), on_stdout=_handle_stdout)
        
        if kdc_found:
            state.inputs['kerberos_config.kdc'] = kdc_found
            state.inputs['kerberos_config.enabled'] = "True"
            if 'Network' not in state.config_data:
                state.config_data['Network'] = {}
            state.config_data['Network']['KDC_AUTO_DETECTED'] = kdc_found
            state.config_data['Network']['KERBEROS_AUTO_ENABLED'] = 'True'
            self.logger.info(f"KDC detected: {kdc_found}")
        
        if progress:
            progress("Network scan complete", 1.0)


class EnvSetupMilestone(Milestone):
    def __init__(self):
        super().__init__("env_setup", "System Environment Setup")
        self.logger = logging.getLogger("EnvSetupMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if progress:
            progress("Setting up environment...", 0.5)
            
        env_keys = ['PROLE_HOME', 'PROLE_CONF', 'PROLE_DATA', 'PROLE_LOGS', 'PROLE_SERVICE', 'NAMESPACE']
        if 'System Environment' not in state.config_data:
            state.config_data['System Environment'] = {}
            
        for k in env_keys:
            val = state.inputs.get(f'env_setup.{k}')
            if val:
                state.config_data['System Environment'][k] = val
        
        if progress:
            progress("Environment setup complete", 1.0)


class SecretManagementMilestone(Milestone):
    def __init__(self):
        super().__init__("secret_management", "Secret Resolution")
        self.logger = logging.getLogger("SecretManagementMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if progress:
            progress("Resolving secrets...", 0.5)

        # Keys that might contain secrets
        secret_keys = [
            'init_password.db_password',
            'init_password.db_password_confirm',
            'kerberos_config.password',
            'init_cluster.k3s_token'
        ]

        for key in secret_keys:
            val = state.inputs.get(key)
            if val:
                resolved = inst_config._resolve_secret_value(val)
                if resolved != val:
                    state.inputs[key] = resolved
                    self.logger.info(f"Resolved secret for {key}")

        # Synchronize confirm password
        if state.inputs.get('init_password.db_password') and not state.inputs.get('init_password.db_password_confirm'):
            state.inputs['init_password.db_password_confirm'] = state.inputs['init_password.db_password']

        if progress:
            progress("Secret resolution complete", 1.0)


class DatabaseCreationMilestone(Milestone):
    def __init__(self):
        super().__init__("init_password", "Database Creation")
        self.logger = logging.getLogger("DatabaseCreationMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if progress:
            progress("Preparing database configuration...", 0.5)
            
        ns = (state.inputs.get('init_password.db_namespace', '') or '').strip()
        if not ns:
            ns = (state.inputs.get('env_setup.NAMESPACE', '') or '').strip() or 'default'
        
        user = (state.inputs.get('init_password.db_username', '') or '').strip()
        p1 = state.inputs.get('init_password.db_password', '')
        
        # In ncurses/silent mode we might not have all fields yet, but we want to ensure defaults
        if not user:
            import getpass
            user = getpass.getuser()
            state.inputs['init_password.db_username'] = user
            
        if not p1:
            import secrets
            import string
            alphabet = string.ascii_letters + string.digits
            p1 = ''.join(secrets.choice(alphabet) for i in range(24))
            state.inputs['init_password.db_password'] = p1
            state.inputs['init_password.db_password_confirm'] = p1
            self.logger.info("Generated default database password.")

        # If password is an unresolvable OpenBao or PROLE_SECRET reference
        # (e.g. clean install where OpenBao hasn't been seeded yet),
        # generate a fresh password so downstream scripts get a real value.
        # OpenBao isn't running yet at this milestone stage, so if the ref
        # can't be resolved right now we must provide a concrete password
        # that will be seeded into OpenBao later by ClusterLifecycleMilestone.
        if p1.startswith('${'):
            resolved = inst_config._resolve_secret_value(p1)
            if resolved and not resolved.startswith('${'):
                p1 = resolved
            else:
                import secrets as _secrets
                import string as _string
                alphabet = _string.ascii_letters + _string.digits
                p1 = ''.join(_secrets.choice(alphabet) for i in range(24))
                self.logger.info("OpenBao/secret ref unresolvable at this stage; generated new DB password for initialization.")
            state.inputs['init_password.db_password'] = p1
            state.inputs['init_password.db_password_confirm'] = p1

        if p1 and not p1.startswith('${') and not inst_config._is_prole_secret(p1):
            if 'Inputs' not in state.config_data:
                state.config_data['Inputs'] = {}
            # Encrypt and save to config data for persistence in prole.cfg
            encrypted_pw = inst_config._encrypt_cfg_secret(p1)
            state.config_data['Inputs']['init_password.db_password'] = encrypted_pw
            state.config_data['Inputs']['init_password.db_password_confirm'] = encrypted_pw
            
        state.inputs['init_password.db_namespace'] = ns
        state.inputs['env_setup.NAMESPACE'] = ns
        
        if 'Database Creation' not in state.config_data:
            state.config_data['Database Creation'] = {}
        state.config_data['Database Creation']['DB_USER'] = user
        state.config_data['Database Creation']['DB_NAME'] = ns
        state.config_data['Database Creation']['NAMESPACE'] = ns
        
        # SSH key generation (simplified for now)
        if self._parse_bool(state.inputs.get('init_password.generate_ssh_key', 'True')):
            key_path = Path.home() / ".ssh" / "id_prole_ed25519"
            if not key_path.exists():
                self._run_cmd(["ssh-keygen", "-t", "ed25519", "-N", "", "-f", str(key_path), "-C", user])

        if progress:
            progress("Database creation preparation complete", 1.0)


class DockerBuildMilestone(Milestone):
    def __init__(self):
        super().__init__("init_db_build", "Docker Build")
        self.logger = logging.getLogger("DockerBuildMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if not self._parse_bool(state.inputs.get('init_db_build.run_build', 'True')):
            self.logger.info("Docker build disabled.")
            return

        if progress:
            progress("Building Prole database image...", 0.1)
            
        tag = state.controller.get_prole_db_version()
        from installer.core.env import _normalize_cluster_env
        env_key = _normalize_cluster_env(state.inputs.get('init_cluster.cluster_env', 'dev')) or 'dev'
        
        args = ["--tag", tag]
        if env_key != 'dev':
            # Add registry prefix for non-dev envs
            server = state.inputs.get('init_cluster.k3s_server_url', 'myrddin.prole.org')
            if '://' in server:
                server = server.split('://')[1]
            if ':' in server:
                server = server.split(':')[0]
            args = ["--tag", f"{server}:5000/prole-db:{tag}", "--push"]
            
        rc = state.controller.run_script("build_db.sh", args=args, on_line=_stream_line)
        if rc != 0:
            self.logger.error(f"Docker build failed (code {rc})")
            
        if progress:
            progress("Docker build complete", 1.0)


class ClusterLifecycleMilestone(Milestone):
    def __init__(self):
        super().__init__("cluster_lifecycle", "Cluster Lifecycle Management")
        self.logger = logging.getLogger("ClusterLifecycleMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        from installer.core.env import _normalize_cluster_env
        cluster_env = state.inputs.get('init_cluster.cluster_env', 'dev')
        env_key = _normalize_cluster_env(cluster_env) or cluster_env
        start_requested = self._parse_bool(state.inputs.get('init_cluster.start_cluster', 'True'))

        if not start_requested:
            self.logger.info("Cluster start not requested.")
            return

        env = self._get_script_env(state)
        
        if env_key == 'dev':
            if progress:
                progress("Initializing k3d dev cluster...", 0.5)
            # Ensure Docker is up (best-effort, macOS auto-start).
            if not state.controller.check_docker_running():
                if sys.platform == "darwin":
                    self.logger.info("Starting Docker...")
                    self._run_cmd(["open", "-a", "Docker"])
                for _ in range(30):
                    time.sleep(2)
                    if state.controller.check_docker_running():
                        break
            if not state.controller.check_docker_running():
                self.logger.error("Docker is not running.")
                return

            cluster_name = "prole-dev-cluster"
            res = subprocess.run(['k3d', 'cluster', 'list', '--no-headers'], capture_output=True, text=True)
            res_stdout = res.stdout or ''
            if cluster_name not in res_stdout:
                from installer.core.env import _k3d_prole_data_volume_args
                prole_data = state.inputs.get('env_setup.PROLE_DATA', str(inst_config.PROJECT_ROOT / "prole-db" / "data"))
                volume_args = _k3d_prole_data_volume_args(prole_data)
                cmd = ['k3d', 'cluster', 'create', cluster_name, '-a', '2'] + volume_args + ['--api-port', '0.0.0.0:6443']
                rc = self._run_cmd(cmd, on_stdout=_stream_line)
                if rc != 0:
                    self.logger.error(f"k3d cluster create failed (code {rc})")
            else:
                if 'running' not in res_stdout.lower():
                    self._run_cmd(['k3d', 'cluster', 'start', cluster_name], on_stdout=_stream_line)

            try:
                subprocess.run(['kubectl', 'config', 'use-context', f'k3d-{cluster_name}'], capture_output=True)
            except Exception:
                pass

            # Ensure prole-db image is available in the cluster (avoid registry pull issues).
            tag = state.controller.get_prole_db_version()
            local_image = f"prole-db:{tag}"
            if subprocess.run(['docker', 'image', 'inspect', local_image], capture_output=True).returncode == 0:
                db_cfg = state.config_data.get('Docker Build', {}) or {}
                registry = (db_cfg.get('LOCAL_REGISTRY_INTERNAL') or '').strip()
                if registry.endswith(".localhost:5000"):
                    registry = registry.replace(".localhost:5000", ":5000")
                elif registry.endswith(".localhost"):
                    registry = registry.replace(".localhost", "")
                if not registry:
                    registry = "k3d-prole-registry:5000"
                remote_tag = f"{registry}/prole-db:{tag}"
                subprocess.run(['docker', 'tag', local_image, remote_tag], capture_output=True)
                self._run_cmd(['k3d', 'image', 'import', remote_tag, '-c', cluster_name], on_stdout=_stream_line)
        
        elif env_key in ('service', 'k3s'):
            if progress:
                progress("Verifying k3s connection...", 0.5)
            
            # Additional logic for remote k3s services could go here
            self.logger.info("Service cluster mode: verification and deployment triggered via scripts.")
            # For now, matching the behavior of silent installer's _step_init_cluster
            # which mostly sets up env and calls init scripts if needed.

        # Initialize OpenBao once cluster is ready. If DB password is anchored,
        # start OpenBao anyway so downstream scripts can attempt to read secrets.
        db_pw_raw = (state.inputs.get('init_password.db_password', '') or '').strip()
        if db_pw_raw:
            db_pw = inst_config._resolve_secret_value(db_pw_raw)
            if db_pw and not db_pw.startswith('${'):
                env_init = dict(env)
                env_init["DB_PASSWORD"] = db_pw
                env_init.setdefault("GRAFANA_ADMIN_PASSWORD", db_pw)
                db_user = (state.inputs.get('init_password.db_username', '') or '').strip()
                if db_user:
                    env_init["PROLE_DB_USER"] = db_user
                rc = state.controller.run_script(
                    "init_openbao.sh",
                    args=["initialize"],
                    env=env_init,
                    stdin_text=f"{db_pw}\n",
                    on_line=_stream_line,
                )
                if rc != 0:
                    self.logger.error(f"OpenBao initialization failed (code {rc})")
            else:
                self.logger.info("OpenBao init skipped: DB password unresolved/anchored. Ensuring OpenBao is running...")
                rc = state.controller.run_script(
                    "init_openbao.sh",
                    args=["start"],
                    env=env,
                    on_line=_stream_line,
                )
                if rc != 0:
                    self.logger.error(f"OpenBao start failed (code {rc})")
        
        if progress:
            progress("Cluster lifecycle management complete", 1.0)


class InitializationScriptsMilestone(Milestone):
    def __init__(self):
        super().__init__("init_scripts", "Initialization Scripts")
        self.logger = logging.getLogger("InitializationScriptsMilestone")

    def _regenerate_port_mapping_cfg(self, state: InstallerState, mode: str) -> None:
        """Regenerate conf/port-mapping.cfg from [Port Forwards] in prole.cfg data."""
        conf_dir = inst_config.PROJECT_ROOT / 'conf'
        mapping_path = conf_dir / 'port-mapping.cfg'
        pf_section = state.config_data.get('Port Forwards', {}) or {}
        prefix = "PORT_FORWARD_K3S_MAPPING_" if mode == "k3s" else "PORT_FORWARD_K3D_MAPPING_"

        # Resolve ${NAMESPACE} from state
        ns = (state.inputs.get('init_password.db_namespace', '') or '').strip()
        if not ns:
            ns = (state.config_data.get('Global', {}) or {}).get('NAMESPACE', '').strip()
        if not ns:
            ns = 'default'

        lines = [
            "# Port mappings for Prole Tools (generated).",
            "# Format: key: local=... remote=... ns=... svc=... address=...",
            ""
        ]
        for k, v in sorted(pf_section.items()):
            if not k.startswith(prefix):
                continue
            parts = {}
            for p in str(v).split(';'):
                if '=' in p:
                    kv = p.split('=', 1)
                    if len(kv) == 2:
                        parts[kv[0].strip()] = kv[1].strip()
            if 'id' not in parts:
                continue
            m_id = parts['id']
            local = parts.get('hostPort', '')
            remote = parts.get('servicePort', '')
            m_ns = parts.get('namespace', '')
            target = parts.get('target', '')
            addr = parts.get('address', '0.0.0.0')
            # Resolve ${NAMESPACE} references
            if m_ns.startswith('${') and 'NAMESPACE' in m_ns:
                m_ns = ns
            if local.startswith('${'):
                continue  # skip entries with unresolvable port refs
            svc = target[4:] if target.startswith('svc/') else target
            lines.append(f"{m_id}: local={local} remote={remote} ns={m_ns} svc={svc} address={addr}")

        try:
            mapping_path.write_text('\n'.join(lines) + '\n')
            self.logger.info(f"Regenerated {mapping_path}")
        except Exception as e:
            self.logger.error(f"Failed to regenerate port-mapping.cfg: {e}")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        env = self._get_script_env(state)
        mode = env.get("PROLE_MODE", "")
        mode_args = ["--mode", mode] if mode else []

        # Regenerate port-mapping.cfg from prole.cfg data to ensure it's clean
        self._regenerate_port_mapping_cfg(state, mode or 'k3d')

        scripts = ["init_common_services.sh", "init_cloudnative_pg.sh"]
        
        # Optional scripts
        if self._parse_bool(state.inputs.get('kerberos_config.enabled', 'False')):
            scripts.append("init_kerberos.sh")
        scripts.extend(["init_prole-db-backup.sh", "init_monitoring.sh", "init_nginx_ingress.sh"])

        for i, script in enumerate(scripts):
            if progress:
                progress(f"Running {script}...", (i / len(scripts)))
            
            # Align args with legacy silent installer behavior
            args = ["update"]
            if script == "init_common_services.sh":
                args = mode_args + ["-n", env.get("SERVICE_NAMESPACE", env.get("NAMESPACE", "default"))]
                if self._parse_bool(state.inputs.get('kerberos_config.enabled', 'False')):
                    args.append("-k")
                args.append("update")
            elif script == "init_cloudnative_pg.sh":
                args = mode_args + ["initialize"]
            elif script == "init_kerberos.sh":
                args = mode_args + ["initialize"]
            elif script == "init_prole-db-backup.sh":
                args = mode_args + ["start"]
            elif script in ("init_monitoring.sh", "init_nginx_ingress.sh"):
                args = mode_args + ["initialize"]
                
            rc = state.controller.run_script(script, args=args, env=env, on_line=_stream_line)
            if rc != 0:
                msg = f"Script {script} failed (code {rc})"
                self.logger.error(msg)
                raise Exception(msg)

        # Verify critical secrets
        ns = env.get("NAMESPACE", "default")
        self.logger.info(f"Verifying critical secrets in namespace {ns}...")
        critical_secrets = ["prole-db-user", "prole-db-superuser", "cnpg-admin-key"]
        missing_secrets = []
        for secret in critical_secrets:
            if self._run_cmd(f"kubectl get secret {secret} -n {ns}") != 0:
                missing_secrets.append(secret)

        if missing_secrets:
            msg = f"CRITICAL: Missing secrets in namespace '{ns}': {', '.join(missing_secrets)}. Database initialization will fail."
            self.logger.error(msg)
            if progress:
                progress(msg, 1.0)
            raise Exception(msg)

        # For k3d: start port-forward daemon and verify Grafana port-forward.
        pf_ok = True
        if mode == "k3d":
            def _pf_line(line: str) -> None:
                self.logger.info(line.rstrip('\n'))

            env_pf = dict(env)
            env_pf.setdefault("PORT_FORWARD_SKIP_VALIDATE", "1")
            env_pf.setdefault("PORT_FORWARD_SKIP_WAIT", "1")
            env_pf.setdefault("PORT_FORWARD_WAIT_TIMEOUT", "20")
            env_pf.setdefault("PORT_FORWARD_WAIT_INTERVAL", "2")

            rc_pf = state.controller.run_script(
                "init_port_forwards.sh",
                args=["--mode", mode, "start"],
                env=env_pf,
                on_line=_stream_line,
            )
            if rc_pf != 0:
                pf_ok = False
                self.logger.error(f"Port-forwards start failed (code {rc_pf})")
            else:
                for _ in range(10):
                    status_lines: list[str] = []

                    def _pf_status_line(line: str) -> None:
                        status_lines.append(line)
                        _stream_line(line)

                    rc_status = state.controller.run_script(
                        "init_port_forwards.sh",
                        args=["--mode", mode, "status", "grafana"],
                        env=env_pf,
                        on_line=_pf_status_line,
                    )
                    running = any(
                        l.lstrip().startswith("RUNNING") and "grafana" in l for l in status_lines
                    )
                    if rc_status == 0 and running:
                        break
                    time.sleep(1)
                else:
                    pf_ok = False
                    self.logger.error("Grafana port-forward not running after start.")

            pf_section = state.config_data.setdefault("Port Forwards", {})
            pf_section["STATUS"] = "Running" if pf_ok else "Failed"
        
        if progress:
            progress("Initialization scripts complete", 1.0)


class DeploymentMilestone(Milestone):
    def __init__(self):
        super().__init__("deployment", "Prole Deployment")
        self.logger = logging.getLogger("DeploymentMilestone")

    def execute(self, state: InstallerState, progress: ProgressCallback | None = None) -> None:
        if progress:
            progress("Deploying Prole...", 0.5)
            
        env = self._get_script_env(state)
        mode = env.get("PROLE_MODE", "")
        mode_args = ["--mode", mode] if mode else []
        # Using init_cloudnative_pg.sh as the primary deployment mechanism
        rc = state.controller.run_script(
            "init_cloudnative_pg.sh",
            args=mode_args + ["deploy", "latest"],
            env=env,
            on_line=_stream_line,
        )
        
        if rc == 0:
            self.logger.info("Deployment successful")
        else:
            msg = f"Deployment failed (code {rc})"
            self.logger.error(msg)
            raise Exception(msg)
            
        if progress:
            progress("Deployment complete", 1.0)