from typing import List, TYPE_CHECKING
from rapidfuzz import fuzz
from pydantic import BaseModel
import pandas as pd
import pathlib
import re
import logging
from itertools import combinations

if TYPE_CHECKING:
    from metadata_extraction import DatasetMeta

logger = logging.getLogger(__name__)

fk_name_pattern = re.compile(r"(_id|_fk|_key|_ref)$", re.IGNORECASE)
pk_name_pattern   = re.compile(r"(^id$|^pk$|_id$|_key$|_pk$)", re.IGNORECASE)

class Relationship(BaseModel):
    source_table: str
    source_column: str
    target_table: str
    target_column: str
    relationship_type: str
    confidence: float
    confidence_name: float
    confidence_values: float

def normalize_col_name(col: str) -> str:
    """Strips surrounding quotes from column names."""
    if not isinstance(col, str):
        return col
    return col.strip().strip('"\'`')

def _is_pk_candidate(col_meta, fk_col_names: set) -> bool:
    col = normalize_col_name(col_meta.column_name)
    if col in fk_col_names:
        return False
    if col_meta.null_count > 0:
        return False
    if col_meta.data_type in ("float", "boolean", "datetime"):
        return False
    # use semantic_type to exclude obvious non-PK columns
    non_pk = {"email", "url", "phone", "address", "description", "currency", "percentage"}
    if col_meta.semantic_type in non_pk:
        return False
    # Use unique_count and total_rows from metadata_extraction
    unique_ratio = col_meta.unique_count / max(1, col_meta.total_rows - col_meta.null_count)
    if unique_ratio < 0.95:
        return False
    if not pk_name_pattern.search(col):
        return False
    return True

def _infer_fk_candidates(dataset: "DatasetMeta") -> dict[str, set[str]]:
    """
    Returns {table_name: {col_name, ...}} of columns that look like FKs by name, so PK detection can skip them.
    """
    return {
        table_name: {
            normalize_col_name(col)
            for col in table_meta.columns
            if fk_name_pattern.search(normalize_col_name(col))
        }
        for table_name, table_meta in dataset.tables.items()
    }
from itertools import combinations

def detect_composite_primary_keys(
    df: pd.DataFrame,
    candidate_cols: list[str],
    min_unique_ratio: float = 0.95,
    max_composite_size: int = 3,
) -> list[tuple[str, ...]]:
    """Given a list of candidate columns that individually failed PK detection, finds combinations that together form a unique key."""
    composite_pks = []
    already_used = set()
    for size in range(2, max_composite_size + 1):
        for combo in combinations(candidate_cols, size):
            if any(col in already_used for col in combo):
                continue
            cols = list(combo)
            if df[cols].isnull().any(axis=None):
                continue
            combined = df[cols].astype(str).agg("-".join, axis=1)
            unique_ratio = combined.nunique() / max(1, len(combined))

            if unique_ratio >= min_unique_ratio:
                composite_pks.append(tuple(cols))
                already_used.update(cols)
                break 
    return composite_pks

def mark_primary_keys(dataset: "DatasetMeta") -> None:
    """
    Two-pass PK marking:
       1. find PKs without FK exclusion
       2. If a non-FK PK exists, exclude FK-named columns
       3. If no non-FK PK found, fall back to all candidates
       4. If still no single PK, try composite keys among non-FK candidates
    """
    fk_candidates = _infer_fk_candidates(dataset)
    for table_name, table_meta in dataset.tables.items():
        df = table_meta.dataframe
        # 1.
        all_pks = [
            col for col, col_meta in table_meta.columns.items()
            if _is_pk_candidate(col_meta, fk_col_names=set())
        ]
        # 2.
        non_fk_pks = [p for p in all_pks if p not in fk_candidates[table_name]]
        #3.
        final_single_pks = non_fk_pks if non_fk_pks else all_pks
        if final_single_pks:
            for col in final_single_pks:
                table_meta.columns[col].is_primary_key = True
            logger.info("Table '%s' PK(s): %s", table_name, final_single_pks)

        else:
            logger.info(
                "Table '%s' no single-column PK found, trying composite...",
                table_name
            )
            composite_candidates = [
                col for col, col_meta in table_meta.columns.items()
                if not pd.api.types.is_float_dtype(df[col])
                and not pd.api.types.is_bool_dtype(df[col])
                and not pd.api.types.is_datetime64_any_dtype(df[col])
                and col_meta.null_count == 0
                and col_meta.unique_count >= 2
            ]
            composite_pks = detect_composite_primary_keys(
                df=df,
                candidate_cols=composite_candidates,
            )
            if composite_pks:
                for combo in composite_pks:
                    for col in combo:
                        table_meta.columns[col].is_primary_key = True
                        table_meta.columns[col].composite_key_with = [
                            c for c in combo if c != col
                        ]
                logger.info("Table '%s' composite PK(s): %s", table_name, composite_pks)
            else:
                logger.warning("Table '%s' no PK detected (single or composite)", table_name)

def _containment_score(fk_vals: set, pk_vals: set) -> float:
    """Fraction of FK values that exist in the PK set."""
    if not fk_vals:
        return 0.0
    return len(fk_vals & pk_vals) / len(fk_vals)

def _fk_name_score(s_col: str, t_col: str, t_table: str) -> float:
    """Compares FK column name against PK column name and target table name."""
    t_col  = normalize_col_name(t_col)
    s_col  = normalize_col_name(s_col)
    s_stripped   = fk_name_pattern.sub("", s_col).lower()
    t_stripped   = fk_name_pattern.sub("", t_col).lower()
    t_table_norm = t_table.lower().rstrip("s")
    direct_sim   = fuzz.ratio(s_col.lower(), t_col.lower()) / 100
    stripped_sim = fuzz.ratio(s_stripped, t_stripped) / 100
    table_sim    = fuzz.ratio(s_stripped, t_table_norm) / 100
    return max(direct_sim, stripped_sim, table_sim)

def _skip_as_fk(
    s_col: str,
    s_meta,
    s_table: str,
    pk_values_per_table: dict,
) -> bool:
    """
    A column should be skipped as FK source only if it's a single PK
    AND there is no other table whose PK it could realistically point to.
    Composite PK members are always scanned as potential FK sources.
    """
    if s_meta.is_primary_key and s_meta.composite_key_with:
        return False
    # skip single PK unless another table has a PK with the same name
    if s_meta.is_primary_key and not s_meta.composite_key_with:
        for t_table, t_pks in pk_values_per_table.items():
            if t_table == s_table:
                continue
            if s_col in t_pks:  # same column name exists as PK elsewhere
                return False    # this might be a FK too
        return True  # no matching PK found elsewhere, skip
    return False

def detect_foreign_keys(
    dataset: "DatasetMeta",
    output_dir: pathlib.Path,
    min_confidence: float = 0.7,
    name_weight: float = 0.5,
    value_weight: float = 0.5,
    pk_sample_cap: int = 10_000,
) -> List[Relationship]:

    output_dir.mkdir(parents=True, exist_ok=True)
    relationships: List[Relationship] = []
    pk_values_per_table: dict[str, dict[str, set]] = {}
    for table_name, table_meta in dataset.tables.items():
        df = table_meta.dataframe
        pk_values_per_table[table_name] = {
            col: set(df[col].dropna().unique()[:pk_sample_cap].tolist())
            for col, meta in table_meta.columns.items()
            if meta.is_primary_key
        }
    for s_table, s_table_meta in dataset.tables.items():
        s_df = s_table_meta.dataframe
        for s_col, s_meta in s_table_meta.columns.items():
            if _skip_as_fk(s_col, s_meta, s_table, pk_values_per_table):
                continue
            if not fk_name_pattern.search(s_col):
                continue
            if s_meta.unique_count < 2: #unlikely to be FK if only 1 unique value
                continue
            s_values = set(
                s_df[s_col].dropna().unique()[:pk_sample_cap].tolist()
            )
            if not s_values:
                continue
            for t_table, t_pks in pk_values_per_table.items():
                if t_table == s_table:
                    continue
                for t_col, t_vals in t_pks.items():
                    if not t_vals:
                        continue
                    name_sim   = _fk_name_score(s_col, t_col, t_table)
                    value_sim  = _containment_score(s_values, t_vals)
                    final_conf = name_weight * name_sim + value_weight * value_sim
                    if final_conf >= min_confidence:
                        relationships.append(
                            Relationship(
                                source_table=s_table,
                                source_column=s_col,
                                target_table=t_table,
                                target_column=t_col,
                                relationship_type="foreign_key",
                                confidence=float(final_conf),
                                confidence_name=float(name_sim),
                                confidence_values=float(value_sim),
                            )
                        )

    _write_markdown(dataset, relationships, output_dir)
    return relationships

def _write_markdown(
    dataset: "DatasetMeta",
    relationships: List[Relationship],
    output_dir: pathlib.Path,
) -> None:
    md_path = output_dir / "list_of_relationships.md"
    with md_path.open("w", encoding="utf-8") as f:
        f.write("# List of PK/FK relationships\n\n")
        f.write(f"Dataset: **{dataset.dataset_name}**\n\n")

        if not relationships:
            f.write("No foreign key relationships detected.\n")
        else:
            for rel in relationships:
                f.write(
                    f"- **{rel.source_table}.{rel.source_column} → "
                    f"{rel.target_table}.{rel.target_column}**  "
                    f"(confidence: {rel.confidence:.3f}, "
                    f"name: {rel.confidence_name:.3f}, "
                    f"values: {rel.confidence_values:.3f})\n"
                )
    logger.info("Markdown saved to: %s", md_path.absolute())

def add_relationships_to_metadata(
    dataset: "DatasetMeta",
    relationships: List[Relationship],
) -> None:
    for rel in relationships:
        col_meta = dataset.tables[rel.source_table].columns[rel.source_column]
        col_meta.foreign_key_to = f"{rel.target_table}.{rel.target_column}"

def detect_primary_and_foreign_keys(
    dataset: "DatasetMeta",
    output_dir: pathlib.Path,
) -> List[Relationship]:
    mark_primary_keys(dataset)
    relationships = detect_foreign_keys(dataset, output_dir=output_dir)
    add_relationships_to_metadata(dataset, relationships)
    return relationships