import yaml
import os
import re
from typing import Set, Dict, Any, Optional, Mapping, Iterable, List
from sqlalchemy import create_engine, inspect
from rich.console import Console

# Import the ConfigurationEngine
from .config_engine import ConfigurationEngine
from sibi_flux.utils.path_utils import PathUtils
from sibi_flux.core import SecureResource, ProjectService

console = Console()


class DiscoveryOrchestrator(SecureResource):
    def __init__(
        self,
        params: Mapping[str, Any],
        rules_path: Optional[str] = "discovery_rules.yaml",
        whitelist_path: str = "discovery_whitelist.yaml",
        registry_path: str = "datacube_registry.yaml",
        db_connection_str: Optional[str] = None,
        field_registry: Optional[Any] = None,
        db_config: Optional[dict] = None,
        all_tables_path: Optional[str] = None,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self.conn_obj = (
            (db_config.get("connection_ref") or db_config.get("connection_obj"))
            if db_config
            else None
        )
        self.config_engine = ConfigurationEngine(
            params, rules_path, context_key=self.conn_obj
        )

        self.whitelist_path = whitelist_path
        self.registry_path = registry_path
        self.db_url = db_connection_str
        self.field_registry = field_registry
        self.db_config = db_config
        self.all_tables_path = all_tables_path
        self.raw_tables: Set[str] = set()

    def run(self, dry_run: bool = False, prune: bool = False):
        """
        Executes the full discovery workflow (Discover + Save).
        Maintained for backward compatibility or single-run calls.
        """
        entries = self.discover()
        self.save_registry(entries, dry_run, prune)

    def discover(self) -> Mapping[str, Any]:
        """
        Executes Steps 1-3 (Read -> Whitelist -> Blacklist -> Rules).
        Returns the dictionary of discovered registry entries.
        """
        db_name = (
            self.db_config.get("id") or self.db_config.get("name") or "default"
            if self.db_config
            else "default"
        )
        console.rule(f"[bold blue]Step 1: Read Schema ({db_name})[/]")
        db_tables = self._step_1_read_schema()
        self.raw_tables = db_tables

        console.rule("[bold blue]Step 2: Apply Whitelist[/]")
        whitelisted_map = self._step_2_apply_whitelist(db_tables)

        # console.rule("[bold blue]Step 2b: Apply Blacklist[/]")
        # Blacklist logic removed

        console.rule("[bold blue]Step 3: Apply Discovery Rules[/]")
        return self._step_3_apply_rules(whitelisted_map)

        # ... (Step 1 and Step 2 omitted for brevity) ...

        # --- Step 2b: Apply Blacklist ---
        # Step 2b Removed
        """
        Removes tables that match any pattern in the blacklist.
        """
        if not self.blacklist_path or not self.get_secure_path(self.blacklist_path).exists():
            console.print("[dim]No blacklist file configured or found. Skipping.[/dim]")
            return table_map

        with self.open_secure(self.blacklist_path, "r") as f:
            data = yaml.safe_load(f) or {}

        blacklist_rules = data.get("blacklist", [])
        if not blacklist_rules:
            console.print("[dim]Blacklist is empty. Skipping.[/dim]")
            return table_map

        filtered_map = {}
        blacklisted_count = 0

        # Pre-compile patterns for performance in loops
        compiled_rules = []
        for rule in blacklist_rules:
            pattern = rule.get("pattern")
            match_type = rule.get("match_type", "exact")
            if match_type == "regex" and pattern:
                compiled_rules.append((re.compile(pattern), match_type))
            else:
                compiled_rules.append((pattern, match_type))

        for table, meta in table_map.items():
            is_blacklisted = False
            for pattern, match_type in compiled_rules:
                if match_type == "exact" and table == pattern:
                    is_blacklisted = True
                    break
                elif match_type == "prefix" and table.startswith(pattern):
                    is_blacklisted = True
                    break
                elif match_type == "regex":
                    if pattern.search(table):
                        is_blacklisted = True
                        break

            if is_blacklisted:
                blacklisted_count += 1
            else:
                filtered_map[table] = meta

        if blacklisted_count > 0:
            console.print(
                f"[yellow]Excluded {blacklisted_count} tables based on blacklist rules.[/yellow]"
            )

        console.print(f"Proceeding with [green]{len(filtered_map)}[/] tables.")
        return filtered_map

    def _step_1_read_schema(self) -> Set[str]:
        """
        Gets table list from all_tables.yaml (preferred) or introspection.
        """
        # OPTIMIZATION: Check all_tables.yaml first
        if self.all_tables_path and self.get_secure_path(self.all_tables_path).exists():
            try:
                with self.open_secure(self.all_tables_path, "r") as f:
                    data = yaml.safe_load(f) or {}

                # Locate our connection's tables
                conn_key = (
                    self.db_config.get("connection_obj")
                    or self.db_config.get("connection_ref")
                )
                if conn_key and conn_key in data:
                     tables = set(data[conn_key])
                     console.print(f"Loaded [green]{len(tables)}[/] tables from cached inventory.")
                     return tables
            except Exception as e:
                console.print(f"[yellow]Failed to read cached inventory: {e}. Falling back to DB.[/yellow]")

        # Fallback to DB Introspection
        """Introspects the database to get the raw list of actual tables."""
        if not self.db_url:
            raise ValueError("No database connection string provided and no cached tables found.")

        try:
            engine = create_engine(self.db_url)
            inspector = inspect(engine)
            tables = set(inspector.get_table_names())
            console.print(f"Found [green]{len(tables)}[/] tables in the database.")
            return tables
        except Exception as e:
            console.print(f"[red]Error connecting to DB: {e}[/red]")
            raise

    # --- Step 2: Read Whitelist ---
    def _step_2_apply_whitelist(self, db_tables: Set[str]) -> Dict[str, Any]:
        """
        Intersects the DB tables with the whitelist.
        Returns a map {table_name: whitelist_metadata}.
        If metadata is missing (legacy whitelist), value is {}.
        """
        if not self.get_secure_path(self.whitelist_path).exists():
            console.print(
                "[dim]No whitelist file found. Proceeding in Rule-Driven Mode (using configurator exclusions).[/dim]"
            )
            return {t: {} for t in db_tables}

        with self.open_secure(self.whitelist_path, "r") as f:
            data = yaml.safe_load(f) or {}

        # 1. Scoped Lookup (New Standard)
        key = self.config_engine.context_key
        raw = data
        whitelist_globals = {}

        if key and isinstance(data, dict) and key in data:
            raw = data[key]
            # Capture Globals from Scope before unwrapping tables
            if isinstance(raw, dict):
                if "field_maps_dir" in raw:
                    whitelist_globals["field_maps_dir"] = raw["field_maps_dir"]
                if "datacubes_dir" in raw:
                    whitelist_globals["datacubes_dir"] = raw["datacubes_dir"]

            # Unwrap 'tables' if nested
            if isinstance(raw, dict) and "tables" in raw:
                raw = raw["tables"]

        # 2. Legacy/Global Fallback
        elif isinstance(data, dict) and "tables" in data:
            raw = data["tables"]
            # Try to capture globals from root if valid? Less common in new structure.
            if "field_maps_dir" in data:
                whitelist_globals["field_maps_dir"] = data["field_maps_dir"]

        whitelisted_map = {}
        if isinstance(raw, dict):
            whitelisted_map = raw
        elif isinstance(raw, list):
            whitelisted_map = {t: {} for t in raw}

        if not whitelisted_map:
            console.print(
                "[yellow]Whitelist found but empty. Processing ALL tables.[/]"
            )
            return {t: {} for t in db_tables}

        # 1. Validation: Whitelisted items missing from DB
        wl_keys = set(whitelisted_map.keys())
        missing = wl_keys - db_tables
        if missing:
            console.print(
                f"[red]Warning: {len(missing)} tables in whitelist not found in DB:[/]"
            )
            for t in missing:
                console.print(f"  - {t}")

        # 2. Filtering: Only process tables that exist in both
        valid_keys = db_tables.intersection(wl_keys)

        final_map = {}
        for k in valid_keys:
            item_meta = whitelisted_map[k].copy()
            # Inject Globals
            item_meta.update(whitelist_globals)
            final_map[k] = item_meta

        console.print(f"Proceeding with [green]{len(final_map)}[/] whitelisted tables.")

        return final_map

    # --- Step 3: Apply Rules ---
    def _step_3_apply_rules(self, table_map: Dict[str, Any]) -> Mapping[str, Any]:
        """
        Runs the ConfigurationEngine against the filtered table map.
        Prioritizes metadata from whitelist (e.g. 'path') over rules.
        """
        registry_updates = {}
        tables = set(table_map.keys())

        # Optional: Register columns if field_registry is present
        if self.field_registry and self.db_url:
            try:
                engine = create_engine(self.db_url)
                inspector = inspect(engine)
                console.print(
                    f"Registering fields for {len(tables)} tables in Global Registry..."
                )

                for table in tables:
                    try:
                        cols = inspector.get_columns(table)
                        for col in cols:
                            # Register with default logic (creates TODO if new)
                            self.field_registry.register_field(
                                original_name=col["name"]
                            )
                    except Exception as e:
                        # Log but don't fail discovery
                        pass

            except Exception as e:
                console.print(f"[red]Error during field registration: {e}[/red]")

        for table, meta in table_map.items():
            # OPTIMIZATION: If whitelist provided datacube_path & field_map_path (compiled), SKIP ConfigurationEngine
            if meta and "datacube_path" in meta and "field_map_path" in meta:
                 # Fully defined in whitelist (Compiled mode)
                 final_entry = meta.copy()
                 final_entry["path"] = meta["datacube_path"] # Normalize key
                 # Normalize field map path to module string
                 fm_raw = meta["field_map_path"]
                 if isinstance(fm_raw, str):
                    final_entry["field_map"] = fm_raw.replace("/", ".").replace(".py", "").strip(".")
                 else:
                    final_entry["field_map"] = str(fm_raw)
                 
                 # Ensure .field_map suffix
                 if not final_entry["field_map"].endswith(".field_map"):
                     final_entry["field_map"] += ".field_map"

                 # Defaults from config/meta if missing
                 if not final_entry.get("class_name"):
                      # Fallback calc
                      suffix = self.config_engine.params.get("defaults", {}).get("class_suffix", "Dc")
                      final_entry["class_name"] = "".join(x.title() for x in table.split("_")) + suffix

                 registry_updates[table] = {
                     "path": final_entry["path"],
                     "field_map": final_entry["field_map"],
                     "connection_obj": self.conn_obj,
                     "class_name": final_entry["class_name"],
                     "custom_name": final_entry.get("custom_name")
                 }
                 continue

            # delegated to the Logic Engine
            match_result = self.config_engine.resolve_table(
                table, db_config=self.db_config
            )

            if match_result:
                # Merge Whitelist Meta (Priority) > Rule Match Result
                final_entry = match_result.copy()
                if meta:
                    # Update fields if present in whitelist meta
                    # Common overrides: path, field_map, class_name
                    # Also respect domain and output_template

                    # If whitelist overrides output_template, we must recalculate the path.
                    if "output_template" in meta:
                        # 1. Recalculate Path
                        # We need the 'datacubes_dir' from config (which config_engine uses)
                        cubes_root = self.config_engine.config.settings.cubes_root_path
                        from pathlib import Path

                        # Simple join for now, assuming template is safe
                        final_entry["path"] = str(
                            Path(cubes_root) / meta["output_template"]
                        )

                    # 1b. Explicit Path Override (New Standard)
                    if "datacube_path" in meta:
                        try:
                            # Ensure absolute path (whitelist has project-relative)
                            final_entry["path"] = str(
                                Path.cwd() / meta["datacube_path"]
                            )
                        except:
                            final_entry["path"] = meta["datacube_path"]

                    # 1c. Explicit Field Map Override
                    if "field_map_path" in meta:
                        try:
                            fm_str = PathUtils.to_module_path(meta["field_map_path"])
                            if not fm_str.endswith(".field_map"):
                                fm_str += ".field_map"
                            final_entry["field_map"] = fm_str
                        except:
                            fm_str = meta["field_map_path"]
                            if not fm_str.endswith(".field_map"):
                                fm_str += ".field_map"
                            final_entry["field_map"] = fm_str

                    # Regular merges
                    for k in [
                        "path",
                        "field_map",
                        "connection_obj",
                        "class_name",
                        "custom_name",
                        "domain",
                    ]:
                        if k in meta:
                            final_entry[k] = meta[k]

                # Priority: custom_name > class_name
                if final_entry.get("custom_name"):
                    final_entry["class_name"] = final_entry["custom_name"]

                # Re-validate that path exists or is valid? No, just trust whitelist + template logic.

                registry_updates[table] = {
                    "path": final_entry["path"],
                    "field_map": final_entry["field_map"],
                    "connection_obj": final_entry.get("connection_obj") or final_entry.get("connection_ref"),
                    "class_name": final_entry.get(
                        "class_name"
                    ),  # Preserve class_name if exists
                }
                # Always include custom_name, defaulting to None
                registry_updates[table]["custom_name"] = final_entry.get("custom_name")
            else:
                # OPTIMIZATION: If no rule match, but whitelist HAS "output_template" and "domain", we can construct it!
                # This supports JIT discovery completely driven by whitelist without rules for specific tables.

                # console.print(f"DEBUG JIT: {table} -> Meta Keys: {list(meta.keys()) if meta else 'None'}")

                if meta and "output_template" in meta and "domain" in meta:
                    # Construct Manual Entry
                    cubes_root = self.config_engine.config.settings.cubes_root_path
                    from pathlib import Path

                    full_path = str(Path(cubes_root) / meta["output_template"])

                    # Infer Field Map Path (Best Effort)
                    # Logic: {base_pkg}.{db_domain}.{fields_root}.{domain}.{table}.field_map

                    # Resolving base pkg from params (accessed via self.config_engine.params if available?)
                    # self.config_engine.params is available.

                    params = self.config_engine.params
                    raw_prefix = (
                        params.get("paths", {})
                        .get("target", {})
                        .get("datacubes_dir", "solutions/dataobjects/gencubes/")
                    )

                    # Relativize to support correct module path generation
                    try:
                        prj_root = Path.cwd()
                        prefix_path = Path(raw_prefix)
                        if prefix_path.is_absolute():
                            try:
                                # Try strictly relative to CWD
                                rel_path = prefix_path.relative_to(prj_root)
                            except ValueError:
                                # Fallback: simple string heuristic if path not relative to CWD?
                                # e.g. path is absolute but symlinked?
                                # Just try to grab 'dataobjects' onwards
                                parts = prefix_path.parts
                                if "dataobjects" in parts:
                                    idx = parts.index("dataobjects")
                                    rel_path = Path(*parts[idx:])
                                else:
                                    # Last resort logic
                                    rel_path = prefix_path
                        else:
                            rel_path = prefix_path

                        folder_prefix = str(rel_path)
                    except Exception:
                        folder_prefix = raw_prefix

                    base_pkg = folder_prefix.replace("/", ".").strip(".")
                    if base_pkg.startswith("."):
                        base_pkg = base_pkg[1:]

                    db_domain = "common"
                    if self.db_config:
                        db_domain = self.db_config.get("db_domain", "common")

                    # 1. Datacube Path (Prefer explicit whitelist path)
                    if "datacube_path" in meta:
                        try:
                            # Whitelist paths are relative to project root. Ensure absolute.
                            full_path = str(Path.cwd() / meta["datacube_path"])
                        except:
                            full_path = meta["datacube_path"]
                    else:
                        cubes_root = self.config_engine.config.settings.cubes_root_path
                        from pathlib import Path

                        full_path = str(Path(cubes_root) / meta["output_template"])

                    # 2. Field Map Path (Prefer explicit whitelist path)
                    if "field_map_path" in meta:
                        try:
                            field_map_module = PathUtils.to_module_path(meta["field_map_path"])
                            if not field_map_module.endswith(".field_map"):
                                field_map_module += ".field_map"
                        except:
                            field_map_module = meta["field_map_path"]
                            if not field_map_module.endswith(".field_map"):
                                field_map_module += ".field_map"
                    else:
                        if field_maps_dir:
                            # Sanitize: package-style to filesystem-style
                            if "." in field_maps_dir and "/" not in field_maps_dir and not Path(field_maps_dir).exists():
                                fm_fs_dir = field_maps_dir.replace(".", "/")
                            else:
                                fm_fs_dir = field_maps_dir

                            try:
                                prj_root = Path.cwd()
                                base_fm_pkg = PathUtils.to_module_path(fm_fs_dir, root_path=prj_root)

                                field_map_module = (
                                    f"{base_fm_pkg}.{db_domain}."
                                    f"{meta['domain']}.{table}"
                                )
                            except Exception:
                                base_fm_pkg = field_maps_dir.replace("/", ".")
                                field_map_module = (
                                    f"{base_fm_pkg}.{db_domain}."
                                    f"{meta['domain']}.{table}"
                                )
                        else:
                            # Old Fallback logic
                            params = self.config_engine.params
                            raw_prefix = (
                                params.get("paths", {})
                                .get("target", {})
                                .get("datacubes_dir", "solutions/dataobjects/gencubes/")
                            )
                            try:
                                prj_root = Path.cwd()
                                prefix_path = Path(raw_prefix)
                                if prefix_path.is_absolute():
                                    try:
                                        rel_path = prefix_path.relative_to(prj_root)
                                    except ValueError:
                                        parts = prefix_path.parts
                                        if "dataobjects" in parts:
                                            idx = parts.index("dataobjects")
                                            rel_path = Path(*parts[idx:])
                                        else:
                                            rel_path = prefix_path
                                else:
                                    rel_path = prefix_path
                                folder_prefix = str(rel_path)
                            except Exception:
                                folder_prefix = raw_prefix

                            base_pkg = PathUtils.to_module_path(folder_prefix)
                            if base_pkg.startswith("."):
                                base_pkg = base_pkg[1:]

                            fields_root = (
                                params.get("generation", {}).get("fields_subpackage")
                                or "fields"
                            )
                            field_map_module = (
                                f"{base_pkg}.{db_domain}."
                                f"{fields_root}."
                                f"{meta['domain']}.{table}.field_map"
                            )

                    # Derive Class Name
                    custom_name = meta.get("custom_name")
                    if custom_name:
                        class_name = custom_name
                    else:
                        # Default CamelCase
                        suffix = params.get("defaults", {}).get("class_suffix", "Dc")
                        class_name = (
                            "".join(x.title() for x in table.split("_")) + suffix
                        )

                    # Connection Object
                    conn_obj = meta.get("connection_obj")
                    if not conn_obj and self.db_config:
                        conn_obj = self.db_config.get("connection_obj") or self.db_config.get("connection_ref")

                    registry_updates[table] = {
                        "path": full_path,
                        "field_map": field_map_module,
                        "connection_obj": conn_obj,
                        "class_name": class_name,
                        "custom_name": custom_name,
                    }
                    # console.print(f"[dim]JIT Mapped {table} from whitelist.[/dim]")
                    continue

                console.print(f"[dim]Skipping {table}: No matching discovery rule.[/]")

        console.print(
            f"Mapped [green]{len(registry_updates)}[/] tables to configuration rules."
        )
        return registry_updates

    # --- Step 4: Generate Datacube Registry ---
    def save_registry(
        self,
        new_entries: Mapping[str, Any],
        dry_run: bool,
        prune: bool,
        global_imports: Optional[Iterable[str]] = None,
    ):
        """
        Merges new entries into existing registry and saved.
        """
        console.rule("[bold blue]Step 4: Generate Registry[/]")
        # Load existing
        current_data = {}
        if self.get_secure_path(self.registry_path).exists():
            with self.open_secure(self.registry_path, "r") as f:
                current_data = yaml.safe_load(f) or {}

        if "tables" not in current_data:
            current_data["tables"] = {}

        # Update Global Imports if provided
        if global_imports:
            current_data["global_imports"] = sorted(
                list(set(current_data.get("global_imports", [])).union(global_imports))
            )

        if prune:
            # LEGACY CLEANUP: Rebuild registry to enforce 'tables' pattern
            # This wipes legacy connection-root keys like 'replica_conf'
            clean_data = {}
            
            # 1. Preserve critical headers
            preserve_keys = [
                "global_imports",
                "type_mappings",
                "defaults",
                "base_datacube_class",
                "default_backend",
                "target_datacubes_folder",
                "target_fieldmaps_folder",
                "valid_libpaths",
                "valid_fieldmap_paths"
            ]
            for key in preserve_keys:
                if key in current_data:
                    clean_data[key] = current_data[key]

            # 2. Assign new tables
            clean_data["tables"] = new_entries
            
            # 3. Replace current data
            current_data = clean_data
            
            console.print(
                "[yellow]Prune active: Registry fully rewritten (Legacy keys removed).[/]"
            )
        else:
            # Smart Merge: Preserve 'custom_name' and 'class_name' from existing entries
            # if they are not explicitly overridden by the new entry.
            for table, new_meta in new_entries.items():
                if table in current_data["tables"]:
                    existing = current_data["tables"][table]
                    
                    # 1. Preserve custom_name if new entry doesn't specify one
                    if not new_meta.get("custom_name") and existing.get("custom_name"):
                        new_meta["custom_name"] = existing["custom_name"]
                        
                        # 2. Also preserve class_name (since it's driven by custom_name)
                        # We only preserve class_name if custom_name was preserved 
                        # AND new entry didn't explicitly change class_name logic (unlikely unless configured)
                        if existing.get("class_name"):
                            new_meta["class_name"] = existing["class_name"]

                current_data["tables"][table] = new_meta

        # Sort tables for readability
        current_data["tables"] = dict(sorted(current_data["tables"].items()))

        if dry_run:
            console.print("[yellow]DRY RUN: Registry would be updated with keys:[/]")
            console.print(list(new_entries.keys()))
        else:
            with self.open_secure(self.registry_path, "w") as f:
                yaml.dump(current_data, f, sort_keys=False)
            console.print(f"[bold green]Success! Updated {self.registry_path}[/]")
