"""
MigMan ETL Transform Module.

This module handles the transformation phase of the MigMan ETL pipeline.
It processes the extracted data, applies business rules, data cleaning, and formatting
to prepare the data for loading into the target system.

The transformation process typically includes:
1. Data validation and quality checks
2. Data type conversions and formatting
3. Business rule application
4. Data enrichment and calculated fields
5. Data structure normalization
6. Comprehensive logging throughout the process

Classes:
    MigManTransform: Main class handling MigMan data transformation.
"""

import logging
import re
from typing import Union
from nemo_library_etl.adapter._utils.db_handler_local import ETLDuckDBHandler
from nemo_library_etl.adapter.migman.config_models_migman import (
    ConfigMigMan,
    TransformDuplicatesConfig,
)
from nemo_library_etl.adapter._utils.enums import ETLStep
from nemo_library_etl.adapter._utils.file_handler import ETLFileHandler
from nemo_library import NemoLibrary

from nemo_library_etl.adapter.migman.enums import MigManTransformStep
from rapidfuzz import fuzz


class MigManTransformDuplicate:
    """
    Handles transformation of extracted MigMan data.

    This class manages the transformation phase of the MigMan ETL pipeline,
    providing methods to process, clean, and format the extracted data for loading
    into the target system.

    The transformer:
    - Uses NemoLibrary for core functionality and configuration
    - Integrates with Prefect logging for pipeline visibility
    - Applies business rules and data validation
    - Handles data type conversions and formatting
    - Provides data enrichment and calculated fields
    - Ensures data quality and consistency

    Attributes:
        nl (NemoLibrary): Core Nemo library instance for system integration.
        config: Configuration object from the Nemo library.
        logger: Prefect logger for pipeline execution tracking.
        cfg (PipelineMigMan): Pipeline configuration with transformation settings.
    """

    def __init__(
        self,
        nl: NemoLibrary,
        cfg: ConfigMigMan,
        logger: Union[logging.Logger, object],
        fh: ETLFileHandler,
        local_database: ETLDuckDBHandler,
    ) -> None:
        """
        Initialize the MigMan Transform instance.

        Sets up the transformer with the necessary library instances, configuration,
        and logging capabilities for the transformation process.

        Args:
            nl (NemoLibrary): Core Nemo library instance for system integration.
            cfg (PipelineMigMan): Pipeline configuration object containing
                                                          transformation settings and rules.
            logger (Union[logging.Logger, object]): Logger instance for recording execution details.
                                                   Can be a standard Python logger or Prefect logger.
        """
        self.nl = nl
        self.cfg = cfg
        self.logger = logger
        self.fh = fh
        self.local_database = local_database

        super().__init__()

    def duplicates(self) -> None:
        """
        Handle duplicate records in the MigMan data.

        This method identifies and processes duplicate records in the extracted
        MigMan data to ensure data integrity and quality before loading into the
        target system. The specific logic for handling duplicates should be
        implemented based on business rules and requirements.

        Steps may include:
        1. Identifying duplicate records based on key fields.
        2. Merging or removing duplicates according to defined rules.
        3. Logging actions taken for audit purposes.

        Note:
            The actual implementation of duplicate handling logic is pending
            and should be customized to fit the MigMan system's needs.
        """
        self.logger.info("Handling duplicates in MigMan data")

        if not self.cfg.transform.duplicate.active:
            self.logger.info("Duplicate configuration is inactive, skipping joins")
            return

        # create UDFs in DuckDB
        def normalize_text(s: str) -> str:
            if s is None:
                return ""
            s = s.lower().strip()
            s = (
                s.replace("ä", "ae")
                .replace("ö", "oe")
                .replace("ü", "ue")
                .replace("ß", "ss")
            )
            s = re.sub(r"[^\w\s]", " ", s)  # remove punctuation
            s = re.sub(r"\s+", " ", s).strip()  # collapse spaces
            return s

        def text_similarity(a: str, b: str) -> float:
            a = normalize_text(a)
            b = normalize_text(b)
            return float(fuzz.token_set_ratio(a, b))

        self.local_database.con.create_function("normalize_text", normalize_text)
        self.local_database.con.create_function("text_similarity", text_similarity)

        for duplicate_name, model in self.cfg.transform.duplicate.duplicates.items():
            if model.active is False:
                self.logger.info(f"Skipping inactive duplicate model: {duplicate_name}")
                continue

            if not duplicate_name in self.cfg.setup.projects:
                self.logger.warning(
                    f"Duplicate model {duplicate_name} not in configured projects, skipping"
                )
                continue

            self._perform_duplicate_check(duplicate_name, model)

    def _perform_duplicate_check(
        self, duplicate_name: str, model: TransformDuplicatesConfig
    ) -> None:
        """
        Single-pass duplicate annotation with 100% recall for token_set similarity:
        - Build token inverted index (id -> tokens of normalized text)
        - Candidates = pairs sharing at least one token (no false negatives for token_set)
        - Compute similarity and aggregate partners
        - Output has same row count as source; includes JSON partners, top score, count
        - Logs progress per phase (row/token/candidate counts)
        """
        from datetime import datetime

        start_time = datetime.now()

        table = self.local_database.latest_table_name(
            steps=MigManTransformStep,
            maxstep=MigManTransformStep.DUPLICATES,
            entity=duplicate_name,
        )
        if table is None:
            raise ValueError(f"No table found for entity {duplicate_name}")
        id_col = model.primary_key
        con = self.local_database.con
        thresh = model.threshold
        out_tbl = f"{MigManTransformStep.DUPLICATES.value}_{duplicate_name}"

        self.logger.info(
            f"[{duplicate_name}] Starting duplicate check on table {table} for fields {model.fields} with threshold {thresh}"
        )

        # Helper: CONCAT of configured columns (safe cast to VARCHAR)
        def concat_expr(alias: str, cols: list[str]) -> str:
            parts = [f"coalesce(CAST({alias}.\"{c}\" AS VARCHAR), '')" for c in cols]
            return " || ' | ' || ".join(parts) if parts else "''"

        base_concat = concat_expr("b", model.fields)

        member_concat = concat_expr("m", model.fields)

        # Fill-score: count non-empty values across the fields used for duplicate check
        fill_terms = []
        for c in model.fields:
            fill_terms.append(
                f"CASE WHEN NULLIF(trim(COALESCE(CAST(\"{c}\" AS VARCHAR), '')), '') IS NULL THEN 0 ELSE 1 END"
            )
        fill_score_expr = " + ".join(fill_terms) if fill_terms else "0"
        # 1) Base: id, normalized text, raw length
        base_tmp = f"__dup_base_{duplicate_name.lower()}"
        con.execute(f'DROP TABLE IF EXISTS "{base_tmp}"')
        con.execute(
            f"""
            CREATE TEMPORARY TABLE "{base_tmp}" AS
            SELECT
                CAST(b."{id_col}" AS VARCHAR)                 AS id,
                normalize_text({base_concat})                 AS norm_text,
                length({base_concat})                         AS txt_len
            FROM "{table}" b
        """
        )
        n_rows = con.execute(f'SELECT COUNT(*) FROM "{base_tmp}"').fetchone()[0]
        self.logger.info(
            f"[{duplicate_name}] Phase 1/5: base prepared — rows: {n_rows:,}"
        )

        # (Optional) Tiny stopword set; leave empty list [] for absolute maximal recall.
        # For high threshold (>=90), filtering 'gmbh'/'ag' etc. is usually safe and reduces huge candidate stars.
        stop_tmp = f"__dup_stop_{duplicate_name.lower()}"
        con.execute(f'DROP TABLE IF EXISTS "{stop_tmp}"')
        con.execute(
            f"""
            CREATE TEMPORARY TABLE "{stop_tmp}"(token VARCHAR);
            INSERT INTO "{stop_tmp}" VALUES
                -- comment-out or remove lines to disable specific stopwords
                ('gmbh'),('mbh'),('kg'),('ag'),('co'),('kg'),('und'),('the'),('der'),('die');
        """
        )

        # 2) Tokenize: one row per (id, token)
        # Use simple space split on the normalized text (already lowercased & cleaned).
        tokens_tmp = f"__dup_tokens_{duplicate_name.lower()}"
        con.execute(f'DROP TABLE IF EXISTS "{tokens_tmp}"')
        con.execute(
            f"""
            CREATE TEMPORARY TABLE "{tokens_tmp}" AS
            SELECT
                id,
                u.token
            FROM "{base_tmp}"
            , UNNEST(string_split(norm_text, ' ')) AS u(token)
            WHERE u.token <> ''
            AND NOT EXISTS (SELECT 1 FROM "{stop_tmp}" s WHERE s.token = u.token);
            """
        )
        n_tokens = con.execute(f'SELECT COUNT(*) FROM "{tokens_tmp}"').fetchone()[0]
        n_dist_t = con.execute(
            f'SELECT COUNT(DISTINCT token) FROM "{tokens_tmp}"'
        ).fetchone()[0]
        self.logger.info(
            f"[{duplicate_name}] Phase 2/5: tokens built — rows: {n_tokens:,}, distinct tokens: {n_dist_t:,}"
        )

        # 3) Candidate pairs via token matches (distinct id pairs)
        #pairs_cand = f"__dup_pairs_cand_{duplicate_name.lower()}"
        #con.execute(f'DROP TABLE IF EXISTS "{pairs_cand}"')
        #con.execute(
        #    f"""
        #    CREATE TEMPORARY TABLE "{pairs_cand}" AS
        #    SELECT DISTINCT
        #        LEAST(a.id, b.id)  AS left_id,
        #        GREATEST(a.id, b.id) AS right_id
        #    FROM "{tokens_tmp}" a
        #    JOIN "{tokens_tmp}" b
        #    ON a.token = b.token
        #    AND a.id <> b.id
        #"""
        #)
        pairs_cand = f"__dup_pairs_cand_{duplicate_name.lower()}"
        con.execute(f'DROP TABLE IF EXISTS "{pairs_cand}"')
        con.execute(
            f"""
            CREATE TEMPORARY TABLE "{pairs_cand}" AS
            SELECT
                a.id AS left_id,
                b.id AS right_id
            FROM "{tokens_tmp}" a
            JOIN "{tokens_tmp}" b
              ON a.token = b.token
             AND a.id < b.id
            GROUP BY 1, 2
            HAVING count(*) >= 2 -- optional: only keep pairs sharing 2+ tokens to reduce candidates (tune for recall/performance)
        """
        )
        
        
        n_cand = con.execute(f'SELECT COUNT(*) FROM "{pairs_cand}"').fetchone()[0]
        self.logger.info(
            f"[{duplicate_name}] Phase 3/5: candidates generated — pairs: {n_cand:,}"
        )

        # 4) Score candidates and filter by threshold
        #pairs_keep = f"__dup_pairs_keep_{duplicate_name.lower()}"
        #con.execute(f'DROP TABLE IF EXISTS "{pairs_keep}"')
        #con.execute(
        #    f"""
        #    CREATE TEMPORARY TABLE "{pairs_keep}" AS
        #    SELECT
        #        p.left_id,
        #        p.right_id,
        #        text_similarity(a.norm_text, b.norm_text) AS score
        #    FROM "{pairs_cand}" p
        #    JOIN "{base_tmp}" a ON a.id = p.left_id
        #    JOIN "{base_tmp}" b ON b.id = p.right_id
        #    WHERE text_similarity(a.norm_text, b.norm_text) >= {thresh}
        #"""
        #)
        pairs_keep = f"__dup_pairs_keep_{duplicate_name.lower()}"
        con.execute(f'DROP TABLE IF EXISTS "{pairs_keep}"')
        con.execute(
            f"""
            CREATE TEMPORARY TABLE "{pairs_keep}" AS
            WITH scored AS (
                SELECT
                    p.left_id,
                    p.right_id,
                    text_similarity(a.norm_text, b.norm_text) AS score
                FROM "{pairs_cand}" p
                JOIN "{base_tmp}" a ON a.id = p.left_id
                JOIN "{base_tmp}" b ON b.id = p.right_id
            )
            SELECT
                left_id,
                right_id,
                score
            FROM scored
            WHERE score >= {thresh}
        """
        )
        n_keep = con.execute(f'SELECT COUNT(*) FROM "{pairs_keep}"').fetchone()[0]
        self.logger.info(
            f"[{duplicate_name}] Phase 4/5: scored+filtered — matches >= {thresh}: {n_keep:,}"
        )
        # 4b) Build duplicate groups (connected components) using Union-Find in Python
        # This gives a stable group_id per component (min id), and later we build the group key "id1 | id2 | ..."
        pairs_edges = con.execute(
            f'SELECT left_id, right_id FROM "{pairs_keep}"'
        ).fetchall()
        
        has_duplicates = bool(pairs_edges)
        
        # Early exit: no duplicates found
        if not has_duplicates:
            self.logger.info(
                f"[{duplicate_name}] No duplicate pairs kept — writing output without group key fields."
            )
        else:
            parent: dict[str, str] = {}
             

            def find(x: str) -> str:
                # path compression
                parent.setdefault(x, x)
                while parent[x] != x:
                    parent[x] = parent[parent[x]]
                    x = parent[x]
                return x

            def union(a: str, b: str) -> None:
                ra = find(a)
                rb = find(b)
                if ra == rb:
                    return
                # deterministic: attach larger root under smaller root (string compare ok if ids are numeric strings of same width;
                # for safety with numeric ids, we will still compute group_id as MIN later in SQL)
                if ra < rb:
                    parent[rb] = ra
                else:
                    parent[ra] = rb

            # union all edges
            for a, b in pairs_edges:
                union(str(a), str(b))

            # build mapping id -> component representative
            # collect all ids seen in edges
            ids = set()
            for a, b in pairs_edges:
                ids.add(str(a))
                ids.add(str(b))

            mapping = [(i, find(i)) for i in ids]  # (id, rep)

            # store mapping into DuckDB
            grp_tmp = f'__dup_groups_{duplicate_name.lower()}'
            con.execute(f'DROP TABLE IF EXISTS "{grp_tmp}"')
            con.execute(f'CREATE TEMPORARY TABLE "{grp_tmp}" (id VARCHAR, rep VARCHAR)')
            con.executemany(f'INSERT INTO "{grp_tmp}" (id, rep) VALUES (?, ?)', mapping)

            self.logger.info(
                f"[{duplicate_name}] Grouping built — groups table: {grp_tmp}, ids in groups: {len(mapping):,}"
            )



        # 5) Build partners for each id and write final annotated table
        #con.execute(
        #    f"""
        #    CREATE OR REPLACE TABLE "{out_tbl}" AS
        #    WITH partners AS (
        #        SELECT left_id AS id,  right_id AS partner_id, score FROM "{pairs_keep}"
        #        UNION ALL
        #        SELECT right_id AS id, left_id  AS partner_id, score FROM "{pairs_keep}"
        #    ),
        #    agg AS (
        #        SELECT
        #            id,
        #            to_json(
        #                list(
        #                    struct_pack(partner_id := partner_id, score := score)
        #                    ORDER BY score DESC, partner_id
        #                )
        #            ) AS duplicate_partners_json,
        #            max(score) AS duplicate_top_score,
        #            count(*)   AS duplicate_match_count
        #        FROM partners
        #        GROUP BY id
        #    )
        #    SELECT
        #        b.*,
        #        coalesce(agg.duplicate_partners_json, '[]') AS duplicate_partners_json,
        #        coalesce(agg.duplicate_top_score, 0)        AS duplicate_top_score,
        #        coalesce(agg.duplicate_match_count, 0)      AS duplicate_match_count
        #    FROM "{table}" b
        #    LEFT JOIN agg ON agg.id = CAST(b."{id_col}" AS VARCHAR)
        #"""
        #)
        if not has_duplicates:
            con.execute(
                f"""
                CREATE OR REPLACE TABLE "{out_tbl}" AS
                WITH partners AS (
                    SELECT left_id AS id,  right_id AS partner_id, score FROM "{pairs_keep}"
                    UNION ALL
                    SELECT right_id AS id, left_id  AS partner_id, score FROM "{pairs_keep}"
                ),
                agg AS (
                    SELECT
                        id,
                        to_json(
                            list(
                                struct_pack(partner_id := partner_id, score := score)
                                ORDER BY score DESC, partner_id
                            )
                        ) AS duplicate_partners_json,
                        max(score) AS duplicate_top_score,
                        count(*)   AS duplicate_match_count
                    FROM partners
                    GROUP BY id
                )
                SELECT
                    b.*,
                    coalesce(agg.duplicate_partners_json, '[]') AS duplicate_partners_json,
                    coalesce(agg.duplicate_top_score, 0)        AS duplicate_top_score,
                    coalesce(agg.duplicate_match_count, 0)      AS duplicate_match_count,
                    ''                                          AS duplicate_group_key,
                    ''                                          AS duplicate_group_values,
                    ''                                          AS duplicate_recommended_id
                FROM "{table}" b
                LEFT JOIN agg ON agg.id = CAST(b."{id_col}" AS VARCHAR)
                """
            )
        else:
            grp_tmp = f'__dup_groups_{duplicate_name.lower()}'

            con.execute(
                f"""
                CREATE OR REPLACE TABLE "{out_tbl}" AS
                WITH
                -- existing partner list for transparency
                partners AS (
                    SELECT left_id AS id,  right_id AS partner_id, score FROM "{pairs_keep}"
                    UNION ALL
                    SELECT right_id AS id, left_id  AS partner_id, score FROM "{pairs_keep}"
                ),
                partner_agg AS (
                    SELECT
                        id,
                        to_json(
                            list(
                                struct_pack(partner_id := partner_id, score := score)
                                ORDER BY score DESC, partner_id
                            )
                        ) AS duplicate_partners_json,
                        max(score) AS duplicate_top_score,
                        count(*)   AS duplicate_match_count
                    FROM partners
                    GROUP BY id
                ),

                -- attach group rep
                grp AS (
                    SELECT id, rep
                    FROM "{grp_tmp}"
                ),

                -- members per group (rep); build group key as sorted ids "id1 | id2 | ..."
                grp_key AS (
                    SELECT
                        rep,
                        string_agg(id, ' | ' ORDER BY id) AS duplicate_group_key
                    FROM grp
                    GROUP BY rep
                ),

                -- values per group: concat of configured fields per member, ordered by member id
                --grp_values AS (
                --    SELECT
                --        g.rep,
                --        string_agg({member_concat}, ' | ' ORDER BY CAST(m."{id_col}" AS VARCHAR)) --AS duplicate_group_values
                --    FROM grp g
                --    JOIN "{table}" m
                --      ON CAST(m."{id_col}" AS VARCHAR) = g.id
                --    GROUP BY g.rep
                --),

                -- recommendation per group: max fill score; tie-break smallest id
                grp_reco AS (
                    SELECT rep, id AS duplicate_recommended_id
                    FROM (
                        SELECT
                            g.rep,
                            g.id,
                            ({fill_score_expr}) AS fill_score,
                            row_number() OVER (
                                PARTITION BY g.rep
                                ORDER BY ({fill_score_expr}) DESC, g.id ASC
                            ) AS rn
                        FROM grp g
                        JOIN "{table}" m
                          ON CAST(m."{id_col}" AS VARCHAR) = g.id
                    ) x
                    WHERE rn = 1
                )

                SELECT
                    b.*,
                    coalesce(partner_agg.duplicate_partners_json, '[]') AS duplicate_partners_json,
                    coalesce(partner_agg.duplicate_top_score, 0)        AS duplicate_top_score,
                    coalesce(partner_agg.duplicate_match_count, 0)      AS duplicate_match_count,

                    -- group fields: repeated on every row of the group, empty for non-duplicates
                    coalesce(grp_key.duplicate_group_key, '')           AS duplicate_group_key,
                    --coalesce(grp_values.duplicate_group_values, '')     AS duplicate_group_values,
                    {base_concat}                                       AS duplicate_group_values,
                    coalesce(grp_reco.duplicate_recommended_id, '')     AS duplicate_recommended_id

                FROM "{table}" b
                LEFT JOIN partner_agg
                  ON partner_agg.id = CAST(b."{id_col}" AS VARCHAR)
                LEFT JOIN grp
                  ON grp.id = CAST(b."{id_col}" AS VARCHAR)
                LEFT JOIN grp_key
                  ON grp_key.rep = grp.rep
                --LEFT JOIN grp_values
                --  ON grp_values.rep = grp.rep
                LEFT JOIN grp_reco
                  ON grp_reco.rep = grp.rep
                """
            )

        src_cnt = con.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
        out_cnt = con.execute(f'SELECT COUNT(*) FROM "{out_tbl}"').fetchone()[0]

        end_time = datetime.now()
        self.logger.info(
            f"[{duplicate_name}] Phase 5/5: annotated table created: {out_tbl} — rows: {out_cnt:,} (source {src_cnt:,}), duration: {end_time - start_time}"
        )

        if self.cfg.transform.dump_files:
            self.local_database.export_table(
                table_name=out_tbl,
                fh=self.fh,
                step=ETLStep.TRANSFORM,
                entity=duplicate_name,
                substep=MigManTransformStep.DUPLICATES,
            )

        if self.cfg.transform.load_to_nemo:
            self.local_database.upload_table_to_nemo(
                table_name=out_tbl,
                project_name=f"{self.cfg.transform.nemo_project_prefix}{out_tbl}",
                delete_temp_files=self.cfg.transform.delete_temp_files,
            )
