import yaml
import socket
import http.client
import ssl
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from urllib.parse import urlparse
from collections import defaultdict
from sibi_flux.init.env_engine import EnvFile, EnvSection
from rich.console import Console

console = Console()


class DiscoveryParamsUpdater:
    """
    Updates conf/discovery_params.yaml and associated modular YAMLs
    to include settings found in the environment configuration.
    """

    @staticmethod
    def _is_database(section: EnvSection) -> bool:
        return section.type.upper() in [
            "POSTGRES",
            "MYSQL",
            "CLICKHOUSE",
            "SQLALCHEMY_DATABASE",
        ]

    @staticmethod
    def _is_filesystem(section: EnvSection) -> bool:
        return section.type.upper() in ["S3", "MINIO", "FILESYSTEM"]

    @staticmethod
    def _is_osmnx(section: EnvSection) -> bool:
        return section.type.upper() == "OSMNX"

    @staticmethod
    def _is_api(section: EnvSection) -> bool:
        return section.type.upper() == "API"

    @staticmethod
    def _update_yaml(path: Path, new_data: Dict[str, Any], key: Optional[str] = None) -> bool:
        """
        Helper to update a YAML file. If key is provided, updates only that key.
        Returns True if changes were made.
        """
        data = {}
        if path.exists():
            with open(path, "r") as f:
                data = yaml.safe_load(f) or {}

        changes_made = False
        if key:
            if key not in data or data[key] != new_data:
                data[key] = new_data
                changes_made = True
        else:
            if data != new_data:
                data.update(new_data)
                changes_made = True

        if changes_made:
            path.parent.mkdir(parents=True, exist_ok=True)
            with open(path, "w") as f:
                yaml.dump(data, f, sort_keys=False, default_flow_style=False)
        
        return changes_made

    MAX_HOSTS_SCAN = 50

    @staticmethod
    def _check_is_alive(host: str, port: Optional[int], timeout: float = 1.0, production_mode: bool = False) -> Tuple[bool, Optional[str]]:
        """
        Check if a host is reachable and attempt to identify it.
        Returns (is_alive, discovered_type).
        """
        if not host or not port:
            return False, None
            
        is_alive = False
        discovered_type = None

        # 1. Basic TCP Connection Check
        try:
            # We use a short timeout and immediate close to avoid lingering sockets
            with socket.create_connection((host, port), timeout=timeout) as sock:
                is_alive = True
                
                # 2. Attempt fingerprinting
                # A. Try HTTP/HTTPS fingerprinting
                if port in [80, 443, 8080, 6000, 8181]:
                    try:
                        # Use HTTPS for 443, HTTP for others
                        if port == 443:
                            context = None
                            if not production_mode:
                                # Create unverified context for development
                                context = ssl._create_unverified_context()
                            conn = http.client.HTTPSConnection(host, port, timeout=timeout, context=context)
                        else:
                            conn = http.client.HTTPConnection(host, port, timeout=timeout)
                            
                        conn.request("HEAD", "/")
                        resp = conn.getresponse()
                        server_header = resp.getheader("Server")
                        if server_header:
                            discovered_type = server_header
                        # MinIO specific check
                        if not discovered_type and (resp.getheader("x-amz-request-id") or resp.getheader("x-minio-deployment-id")):
                            discovered_type = "MinIO"
                        conn.close()
                    except:
                        pass
                
                # B. Try banner grab for common protocols if still unknown
                if not discovered_type:
                    try:
                        sock.settimeout(timeout)
                        # Minimal peek to avoid blocking
                        banner = sock.recv(256)
                        if b"PostgreSQL" in banner:
                            discovered_type = "PostgreSQL"
                        elif b"mysql" in banner or b"MySQL" in banner:
                            discovered_type = "MySQL"
                        elif b"Redis" in banner:
                            discovered_type = "Redis"
                        elif b"SSH" in banner:
                            discovered_type = "SSH"
                    except:
                        pass
                        
        except (socket.timeout, socket.error):
            pass

        return is_alive, discovered_type

    @staticmethod
    def update(project_path: Path, env_data: EnvFile, production_mode: bool = False) -> None:
        conf_dir = project_path / "conf"
        params_path = conf_dir / "discovery_params.yaml"
        db_path = conf_dir / "databases.yaml"
        storage_path = conf_dir / "storage.yaml"
        osmnx_path = conf_dir / "osmnx.yaml"
        api_path = conf_dir / "api.yaml"

        if not params_path.exists():
            # Fallback check for legacy location to assist migration if needed
            legacy_path = project_path / "generators" / "datacubes" / "discovery_params.yaml"
            if legacy_path.exists():
                console.print(f"[yellow]Migrating {legacy_path.name} to {params_path}[/yellow]")
                params_path.parent.mkdir(parents=True, exist_ok=True)
                legacy_path.rename(params_path)
            else:
                console.print(
                    f"[yellow]Warning: {params_path} not found. Skipping discovery params update.[/yellow]"
                )
                return

        try:
            # 1. Update databases.yaml
            db_sections = [
                s for s in env_data.sections if DiscoveryParamsUpdater._is_database(s)
            ]
            
            db_list = []
            for section in db_sections:
                db_id = section.name.lower()
                conf_name = f"{db_id}_conf"
                
                db_domain = db_id
                for var in section.vars:
                    if var.key.endswith("_DATABASE") or var.key == "DATABASE":
                        db_domain = var.value
                        break

                db_list.append({
                    "id": db_id,
                    "connection_ref": conf_name,
                    "db_domain": db_domain,
                    "import_spec": {"module": "conf.credentials", "symbol": conf_name},
                })

            if db_list:
                DiscoveryParamsUpdater._update_yaml(db_path, db_list, key="databases")
                console.print(f"[green]Updated {db_path.name} with {len(db_list)} databases.[/green]")
            elif db_path.exists():
                db_path.unlink()
                console.print(f"[yellow]Removed {db_path.name} (no databases found).[/yellow]")

            # 2. Update storage.yaml (Filesystems)
            fs_sections = [
                s for s in env_data.sections if DiscoveryParamsUpdater._is_filesystem(s)
            ]
            if fs_sections:
                fs_list = []
                for section in fs_sections:
                    fs_id = section.name.lower()
                    conf_name = f"{fs_id}_conf"
                    fs_path = ""
                    for var in section.vars:
                        if var.key.endswith("_FS_PATH") or var.key == "FS_PATH":
                            fs_path = var.value
                            break
                    
                    fs_list.append({
                        "id": fs_id,
                        "connection_ref": conf_name,
                        "path": fs_path,
                        "import_spec": {"module": "conf.credentials", "symbol": conf_name},
                    })

                # Merge with existing filesystems to PRESERVE layouts and custom fields
                storage_data = {}
                if storage_path.exists():
                    with open(storage_path, "r") as f:
                        storage_data = yaml.safe_load(f) or {}
                
                existing_fs = {f["id"]: f for f in storage_data.get("filesystems", [])}
                for fs in fs_list:
                    if fs["id"] in existing_fs:
                        # Preserve layout if it exists
                        old_fs = existing_fs[fs["id"]]
                        if "layout" in old_fs:
                            fs["layout"] = old_fs["layout"]
                        # Preserve any other custom fields
                        for k, v in old_fs.items():
                            if k not in fs:
                                fs[k] = v

                DiscoveryParamsUpdater._update_yaml(storage_path, fs_list, key="filesystems")
                console.print(f"[green]Updated {storage_path.name} with {len(fs_list)} filesystems.[/green]")

                # --- 2b. Optional: Default S3 Layout ---
                # Only if using S3/Minio and no layouts exist yet
                has_s3 = any(s.type.upper() in ["S3", "MINIO"] for s in fs_sections)
                if has_s3:
                    # Check if layouts already exist in storage.yaml or discovery_params.yaml
                    storage_data = {}
                    if storage_path.exists():
                        with open(storage_path, "r") as f:
                            storage_data = yaml.safe_load(f) or {}

                    with open(params_path, "r") as f:
                        core_data = yaml.safe_load(f) or {}

                    existing_layouts = storage_data.get("s3_layouts", {}) or core_data.get("s3_layouts", {})
                    if not existing_layouts:
                        default_layouts = {
                            "default_layout": {
                                "default_storage": ["default_depot"]
                            }
                        }
                        DiscoveryParamsUpdater._update_yaml(storage_path, default_layouts, key="s3_layouts")
                        console.print(f"[green]Added default S3 layout to {storage_path.name}[/green]")
                        existing_layouts = default_layouts

                    # Auto-associate the first S3/Minio filesystem with a layout if none are associated
                    has_layout_usage = any("layout" in fs for fs in fs_list)
                    if not has_layout_usage and existing_layouts:
                        layout_name = list(existing_layouts.keys())[0]
                        layout_associated = False
                        for fs in fs_list:
                            fs_index = fs_list.index(fs)
                            orig_section = fs_sections[fs_index]
                            if orig_section.type.upper() in ["S3", "MINIO"]:
                                fs["layout"] = layout_name
                                console.print(f"[green]Associated filesystem '{fs['id']}' with layout '{layout_name}'[/green]")
                                layout_associated = True
                                break

                        if layout_associated:
                            DiscoveryParamsUpdater._update_yaml(storage_path, fs_list, key="filesystems")

                # --- 2c. OSMnx Depot Provisioning ---
                osmnx_detected = any(DiscoveryParamsUpdater._is_osmnx(s) for s in env_data.sections)
                if osmnx_detected and storage_path.exists():
                    with open(storage_path, "r") as f:
                        storage_data = yaml.safe_load(f) or {}
                    
                    layouts = storage_data.get("s3_layouts", {})
                    if "default_layout" in layouts:
                        if "osmnx_data" not in layouts["default_layout"]:
                            layouts["default_layout"]["osmnx_data"] = ["pbf_files"]
                            DiscoveryParamsUpdater._update_yaml(storage_path, layouts, key="s3_layouts")
                            console.print("[green]Added 'osmnx_data' depot to default S3 layout for OSMnx.[/green]")

                # 3. Update osmnx.yaml
                osmnx_sections = [
                    s for s in env_data.sections if DiscoveryParamsUpdater._is_osmnx(s)
                ]
                if osmnx_sections:
                    section = osmnx_sections[0]
                    storage_ref = None
                    fs_ref = "source"
                    for var in section.vars:
                        k = var.key.upper()
                        if k.endswith("_STORAGE_REF") or k == "STORAGE_REF":
                            storage_ref = var.value
                        elif k.endswith("_FS_REF") or k == "FS_REF":
                            fs_ref = var.value

                    # Heuristic: If storage_ref is not explicit, look for 'osmnx_data' across ALL filesystems
                    if not storage_ref:
                        found_ref = False
                        for fs in fs_list:
                            if "layout" in fs:
                                lays = fs["layout"]
                                if isinstance(lays, str):
                                    lays = [lays]
                                
                                for lay_name in lays:
                                    if lay_name in existing_layouts:
                                        layout_def = existing_layouts[lay_name]
                                        for depot, subdirs in layout_def.items():
                                            if "osmnx_data" in subdirs:
                                                storage_ref = "osmnx_data_storage"
                                                fs_ref = fs["id"]
                                                found_ref = True
                                                break
                                        if found_ref:
                                            break
                            if found_ref:
                                break
                        
                        # Ultimate fallback
                        if not storage_ref:
                            storage_ref = "default_depot_storage"

                    DiscoveryParamsUpdater._update_yaml(osmnx_path, {
                        "storage_ref": storage_ref,
                        "fs_ref": fs_ref
                    }, key="osmnx")
                    console.print(f"[green]Updated {osmnx_path.name} (resolved storage_ref={storage_ref}, fs_ref={fs_ref})[/green]")
            elif osmnx_path.exists():
                osmnx_path.unlink()
                console.print(f"[yellow]Removed {osmnx_path.name} (no OSMnx settings found).[/yellow]")

            # 4. Update api.yaml
            api_sections = [
                s for s in env_data.sections if DiscoveryParamsUpdater._is_api(s)
            ]
            if api_sections:
                section = api_sections[0]
                api_data = {
                    "title": "Sibi Toolkit API",
                    "version": "0.1.0",
                    "router_module": "api.routers",
                    "enable_mcp": False
                }
                for var in section.vars:
                    k = var.key.upper()
                    if k.endswith("_TITLE") or k == "TITLE":
                        api_data["title"] = var.value
                    elif k.endswith("_VERSION") or k == "VERSION":
                        api_data["version"] = var.value
                    elif k.endswith("_ROUTER_MODULE") or k == "ROUTER_MODULE":
                        api_data["router_module"] = var.value
                    elif k.endswith("_ENABLE_MCP") or k == "ENABLE_MCP":
                        api_data["enable_mcp"] = var.value.lower() == "true"

                DiscoveryParamsUpdater._update_yaml(api_path, api_data, key="api")
                console.print(f"[green]Updated {api_path.name}[/green]")

            # 5. Infrastructure Host Discovery (hosts.yaml)
            hosts_path = conf_dir / "hosts.yaml"
            hosts_dict = {} # Use dict to avoid duplicates by ID
            scan_count = 0

            def _add_host(item_id, item_type, host, port, extra=None):
                nonlocal scan_count
                if not host: return
                
                if scan_count >= DiscoveryParamsUpdater.MAX_HOSTS_SCAN:
                    return

                # Perform reachability and identification check
                is_alive, discovered_type = DiscoveryParamsUpdater._check_is_alive(host, port, production_mode=production_mode)
                scan_count += 1
                
                entry = {
                    "id": item_id,
                    "type": item_type,
                    "port": port,
                    "is_alive": is_alive
                }
                if discovered_type:
                    entry["discovered_type"] = discovered_type
                    
                if extra: entry.update(extra)
                hosts_dict[item_id] = entry

            def _scan_vars(vars_list, section_id=None, section_type=None):
                for var in vars_list:
                    k = var.key.upper()
                    if k.endswith("_URL") or k.endswith("_ENDPOINT"):
                        parsed = urlparse(var.value)
                        if parsed.hostname:
                            # Heuristic for type
                            svc_type = "service"
                            dialect = None
                            if "DB" in k or "DATABASE" in k or (section_type and section_type in ["POSTGRES", "MYSQL", "CLICKHOUSE"]):
                                svc_type = "database"
                                dialect = parsed.scheme.split("+")[0] if "+" in parsed.scheme else parsed.scheme
                            elif "REDIS" in k or (section_type == "REDIS"):
                                svc_type = "cache"
                            elif "S3" in k or "MINIO" in k or (section_type and section_type in ["S3", "MINIO"]):
                                svc_type = "storage"
                            
                            extra = {"protocol": parsed.scheme}
                            if dialect: extra["dialect"] = dialect
                            
                            item_id = section_id or k.lower()
                            _add_host(item_id, svc_type, parsed.hostname, parsed.port, extra)

            # A. Scan Sections
            for section in env_data.sections:
                section_type = section.type.upper()
                section_id = section.name.lower()
                
                # Specialized scanners first
                if DiscoveryParamsUpdater._is_database(section):
                    url = next((v.value for v in section.vars if v.key.endswith("_URL")), None)
                    if url:
                        _scan_vars([v for v in section.vars if v.key.endswith("_URL")], section_id, section_type)
                    else:
                        host = next((v.value for v in section.vars if v.key.endswith("_HOST")), None)
                        port = next((v.value for v in section.vars if v.key.endswith("_PORT")), None)
                        if host:
                            _add_host(section_id, "database", host, int(port) if port and port.isdigit() else None, 
                                     {"dialect": section_type.lower() if section_type != "SQLALCHEMY_DATABASE" else "unknown"})
                
                elif DiscoveryParamsUpdater._is_filesystem(section):
                    endpoint = next((v.value for v in section.vars if v.key.endswith("_ENDPOINT") or v.key == "ENDPOINT"), None)
                    if endpoint:
                        _scan_vars([v for v in section.vars if v.key.endswith("_ENDPOINT") or v.key == "ENDPOINT"], section_id, section_type)
                
                elif section_type == "REDIS":
                    host = next((v.value for v in section.vars if v.key == "HOST" or v.key.endswith("_HOST")), None)
                    port = next((v.value for v in section.vars if v.key == "PORT" or v.key.endswith("_PORT")), None)
                    if host:
                        _add_host(section_id, "cache", host, int(port) if port and port.isdigit() else 6379)
                
                else:
                    # Scan generic sections for URLs
                    _scan_vars(section.vars, section_id, section_type)

            # B. Scan Orphans
            _scan_vars(env_data.orphan_vars)

            if hosts_dict:
                if scan_count >= DiscoveryParamsUpdater.MAX_HOSTS_SCAN:
                    console.print(f"[yellow]Warning: Infrastructure discovery hit the safety limit of {DiscoveryParamsUpdater.MAX_HOSTS_SCAN} hosts. Some hosts might be missing.[/yellow]")
                
                # 1. Map signatures to entries
                signature_map = defaultdict(list)
                
                for entry in hosts_dict.values():
                    etype = entry.pop("type", "service")
                    group = "service"
                    if etype == "database": group = "databases"
                    elif etype == "storage": group = "storage"
                    
                    port = entry.pop("port", None)
                    ports = [port] if port else []
                    
                    # Create signature
                    host = entry.get("host")
                    dialect = entry.get("dialect", "")
                    protocol = entry.get("protocol", "")
                    signature = (group, host, tuple(ports), dialect, protocol)
                    
                    # Store entry with metadata
                    entry["_group"] = group
                    entry["ports"] = ports
                    signature_map[signature].append(entry)

                # 2. Build final grouped structure with shared naming
                grouped_hosts = {"databases": [], "storage": [], "service": []}
                shared_counters = defaultdict(int)
                
                # Count shared signatures per group to decide on auto-indexing
                shared_per_group = defaultdict(int)
                for sig, entries in signature_map.items():
                    if len(entries) > 1:
                        shared_per_group[sig[0]] += 1

                for sig, entries in signature_map.items():
                    group = sig[0]
                    if len(entries) > 1:
                        shared_counters[group] += 1
                        idx = shared_counters[group]
                        
                        # Identify the discovered type for naming
                        entry = entries[0]
                        dtype = entry.get("discovered_type")
                        
                        if dtype:
                            # Sanitize discovered_type for ID (lowercase, no special chars)
                            import re
                            clean_type = re.sub(r'[^a-zA-Z0-9]', '_', dtype).lower().strip('_')
                            # Reduce multiple underscores
                            clean_type = re.sub(r'_+', '_', clean_type)
                            base_id = f"shared_{clean_type}"
                            base_name = f"Shared {dtype}"
                        else:
                            type_name = group.rstrip("s")
                            if group == "service": type_name = "service"
                            base_id = f"{type_name}_shared"
                            base_name = base_id.replace("_", " ").title()
                        
                        # Ensure uniqueness by adding index if multiple shared hosts exist in this category
                        id_val = f"{base_id}_{idx}" if shared_per_group[group] > 1 else base_id
                        name_val = f"{base_name} {idx}" if shared_per_group[group] > 1 else base_name
                        
                        entry["id"] = id_val
                        entry["name"] = name_val
                    else:
                        entry = entries[0]
                        if "name" not in entry:
                            entry["name"] = entry["id"].capitalize()
                    
                    # Cleanup internal helper
                    entry.pop("_group", None)
                    grouped_hosts[group].append(entry)

                DiscoveryParamsUpdater._update_yaml(hosts_path, grouped_hosts, key="hosts")
                console.print(f"[green]Updated {hosts_path.name} with shared host discovery logic.[/green]")

            # 6. Legacy cleanup in discovery_params.yaml
            with open(params_path, "r") as f:
                core_data = yaml.safe_load(f) or {}
            
            migrated_found = False
            for legacy in ["databases", "filesystems", "osmnx", "api", "s3_layouts"]:
                if legacy in core_data:
                    del core_data[legacy]
                    migrated_found = True
            
            if migrated_found:
                with open(params_path, "w") as f:
                    yaml.dump(core_data, f, sort_keys=False, default_flow_style=False)
                console.print(f"[yellow]Cleaned up legacy sections from {params_path.name}[/yellow]")

        except Exception as e:
            console.print(f"[red]Failed to update configuration files: {e}[/red]")
