import yaml
from pathlib import Path
from typing import List, Dict, Any, Set
from collections import defaultdict
import re
from rich.console import Console

console = Console()


class RuleEngine:
    def __init__(self, all_tables_path: Path, rules_path: Path):
        self.all_tables_path = all_tables_path
        self.rules_path = rules_path
        self.raw_tables: Dict[str, List[str]] = {}
        self.existing_rules: Dict[str, List[Dict[str, Any]]] = {}

    def load(self):
        """Loads all_tables.yaml and discovery_rules.yaml"""
        if self.all_tables_path.exists():
            with open(self.all_tables_path, "r") as f:
                self.raw_tables = yaml.safe_load(f) or {}

        if self.rules_path.exists():
            with open(self.rules_path, "r") as f:
                self.existing_rules = yaml.safe_load(f) or {}

    def _get_covered_tables(
        self, tables: List[str], rules: List[Dict[str, Any]]
    ) -> Set[str]:
        """Identifies tables that match existing rules."""
        covered = set()
        for table in tables:
            for rule in rules:
                pattern = rule.get("pattern")
                match_type = rule.get("match_type", "prefix")

                is_match = False
                if match_type == "exact" and table == pattern:
                    is_match = True
                elif match_type == "prefix" and table.startswith(pattern):
                    is_match = True
                elif match_type == "regex":
                    if re.search(pattern, table):
                        is_match = True

                if is_match:
                    covered.add(table)
                    break
        return covered

    def _cluster_unmapped_tables(self, tables: List[str]) -> Dict[str, List[str]]:
        """
        Clusters tables by common prefixes.
        Heuristic: Split by '_' and take the first token as the candidate prefix.
        Ignores prefixes with fewer than 3 tables.
        """
        clusters = defaultdict(list)
        for table in tables:
            parts = table.split("_")
            if len(parts) > 1:
                prefix = parts[0]
                clusters[prefix].append(table)
            else:
                # Single word table? Treat as exact match candidate later?
                # Or cluster under 'misc'?
                clusters[table].append(table)  # Self-cluster

        # Filter small clusters
        # Heuristic: Pattern must be significant.
        # But 'datas' (1 table) was its own rule in the example.
        # Let's propose rules for ANY cluster >= 2? Or all?
        # User prompt implies "heuristic". Let's stick to >= 3 for prefix, or explicit unique.

        # Actually, let's just group everything and return significant ones.
        return {k: v for k, v in clusters.items() if len(v) >= 1}

    def propose_rules(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        Generates new rules for unmapped tables.
        Returns a dictionary of new rules keyed by connection.
        """
        updates = {}

        for conn, tables in self.raw_tables.items():
            existing = self.existing_rules.get(conn, [])
            covered = self._get_covered_tables(tables, existing)
            unmapped = sorted(list(set(tables) - covered))

            if not unmapped:
                continue

            console.print(f"[{conn}] Analyzing {len(unmapped)} unmapped tables...")

            clusters = self._cluster_unmapped_tables(unmapped)
            new_rules = []

            for prefix, cluster_tables in clusters.items():
                # Heuristic decision logic

                # Case 1: Single item cluster -> Exact Match
                if len(cluster_tables) == 1 and cluster_tables[0] == prefix:
                    table_name = cluster_tables[0]
                    new_rules.append(
                        {
                            "pattern": table_name,
                            "match_type": "exact",
                            "domain": table_name,  # Default domain
                            "output_template": f"{table_name}_cubes.py",
                        }
                    )
                    continue

                # Case 2: Multi-item prefix cluster -> Prefix Match
                # Only if the cluster size is significant (e.g. >= 2)
                # If only 1 item like 'auth_group', we might still want a prefix rule 'auth_'?
                # Or exact?
                # Canonical example has 'parameters_' (prefix) for many tables.
                # It has 'datas' (exact) for 1 table.

                # Heuristic: If prefix matches the full table name (e.g. 'users' split has 1 part 'users'), it's exact.
                # If prefix is a substring (e.g. 'auth' from 'auth_group'), it's prefix.

                rule_type = "prefix"
                pattern = f"{prefix}_"

                # Check coverage of this proposed rule against the cluster
                # (Ideally it covers all of them)

                new_rules.append(
                    {
                        "pattern": pattern,
                        "match_type": "prefix",
                        "domain": prefix,
                        "output_template": f"{prefix}_cubes.py",
                    }
                )

            if new_rules:
                updates[conn] = new_rules

        return updates

    def save_proposal(self, updates: Dict[str, List[Dict[str, Any]]]):
        """Merges and saves updates to discovery_rules.yaml"""
        if not updates:
            console.print("No new rules to propose.")
            return

        # Deep merge
        final_data = self.existing_rules.copy()
        for conn, rules in updates.items():
            if conn not in final_data:
                final_data[conn] = []

            # Avoid duplicate rules (basic check on pattern)
            existing_patterns = {r["pattern"] for r in final_data[conn]}

            count = 0
            for r in rules:
                if r["pattern"] not in existing_patterns:
                    final_data[conn].append(r)
                    count += 1

            if count > 0:
                console.print(f"[{conn}] Added {count} new rules.")

        # Ensure ALL groups are sorted by domain and pattern (User Request)
        for conn in final_data:
            final_data[conn].sort(key=lambda x: (x.get("domain", ""), x.get("pattern", "")))

        with open(self.rules_path, "w") as f:
            yaml.dump(final_data, f, sort_keys=False)

        console.print(f"[green]Updated {self.rules_path}[/green]")
