from typing import Any, Optional
from smolagents.tools import Tool
import os
import json
import io
import pandas

class DataQualityProfilerTool(Tool):
    name = "data_quality_profiler"
    description = "Profiles tabular datasets (CSV, TSV, JSON lines, Markdown tables, or file paths) according to enterprise Data Governance (CDMP\u00ae / DAMA-DMBOK) standards. Computes overall data quality scores, completeness, uniqueness, primary key candidates, statistical distributions, outlier flags, and actionable anomalies for downstream agent tasks."
    inputs = {'data': {'type': 'string', 'description': 'Raw tabular string content (CSV, TSV, JSON, Markdown table) or an existing file path to inspect.'}, 'file_type': {'type': 'string', 'description': "Optional format hint: 'csv', 'tsv', 'json', 'markdown', or 'auto' (default: 'auto').", 'nullable': True}, 'max_sample_rows': {'type': 'integer', 'description': 'Maximum number of rows to sample for profiling if dataset is huge (default: 10000).', 'nullable': True}}
    output_type = "string"

    def _parse_input(self, data_str: str, file_type: Optional[str] = None):
        import io
        import json
        import os
        import pandas as pd

        clean_str = data_str.strip()
        fmt = (file_type or "auto").lower()

        # Check if input is a valid existing file path
        if os.path.exists(clean_str) and os.path.isfile(clean_str):
            ext = os.path.splitext(clean_str)[1].lower()
            if ext in [".csv", ".txt"] or fmt == "csv":
                return pd.read_csv(clean_str)
            elif ext in [".tsv", ".tab"] or fmt == "tsv":
                return pd.read_csv(clean_str, sep="\t")
            elif ext == ".json" or fmt == "json":
                try:
                    return pd.read_json(clean_str)
                except Exception:
                    return pd.read_json(clean_str, lines=True)
            elif ext == ".parquet":
                return pd.read_parquet(clean_str)

        # Handle Markdown table format
        if fmt == "markdown" or ("|" in clean_str and "\n|-" in clean_str):
            lines = [line.strip() for line in clean_str.splitlines() if line.strip().startswith("|")]
            if len(lines) >= 3:
                header = [c.strip() for c in lines[0].split("|")[1:-1]]
                data_rows = []
                for row_line in lines[2:]:
                    cols = [c.strip() for c in row_line.split("|")[1:-1]]
                    if len(cols) == len(header):
                        data_rows.append(cols)
                if data_rows:
                    return pd.DataFrame(data_rows, columns=header)

        # Handle JSON format (array or jsonl)
        if fmt == "json" or clean_str.startswith("[") or clean_str.startswith("{"):
            try:
                data_obj = json.loads(clean_str)
                if isinstance(data_obj, list):
                    return pd.DataFrame(data_obj)
                elif isinstance(data_obj, dict):
                    return pd.DataFrame([data_obj])
            except Exception:
                try:
                    return pd.read_json(io.StringIO(clean_str), lines=True)
                except Exception:
                    pass

        # Handle Delimited (CSV / TSV)
        if fmt == "tsv" or ("\t" in clean_str and "," not in clean_str):
            return pd.read_csv(io.StringIO(clean_str), sep="\t")
        else:
            try:
                return pd.read_csv(io.StringIO(clean_str))
            except Exception:
                return pd.read_csv(io.StringIO(clean_str), sep=None, engine="python")

    def forward(
        self,
        data: str,
        file_type: Optional[str] = "auto",
        max_sample_rows: Optional[int] = 10000,
    ) -> str:
        """Profiles the provided dataset and returns a comprehensive Markdown report."""
        import pandas as pd

        if not data or not data.strip():
            return "Error: Empty dataset provided. Please provide non-empty tabular data."

        try:
            df = self._parse_input(data, file_type)
        except Exception as err:
            return f"Error parsing input data: {str(err)}. Ensure valid CSV, TSV, JSON, or Markdown table format."

        if df.empty:
            return "Error: Parsed dataframe contains zero rows or columns."

        total_rows = len(df)
        total_cols = len(df.columns)
        is_sampled = False

        if max_sample_rows and total_rows > max_sample_rows:
            df = df.sample(n=max_sample_rows, random_state=42)
            is_sampled = True

        analyzed_rows = len(df)

        # 1. Uniqueness & Row Duplication
        duplicate_rows_count = int(df.duplicated().sum())
        duplicate_rate = (duplicate_rows_count / analyzed_rows) if analyzed_rows > 0 else 0.0

        # 2. Completeness Dimension
        total_cells = analyzed_rows * total_cols
        missing_cells = int(df.isna().sum().sum())
        completeness_rate = 1.0 - (missing_cells / total_cells) if total_cells > 0 else 1.0

        # Primary Key Candidates
        pk_candidates = [
            str(col) for col in df.columns
            if df[col].isna().sum() == 0 and df[col].nunique() == analyzed_rows and analyzed_rows > 1
        ]

        # 3. Column Profiling & Anomaly Detection
        anomalies = []
        col_profiles = []

        if duplicate_rows_count > 0:
            anomalies.append(
                f"**High Row Duplication**: Found {duplicate_rows_count} exact duplicate rows ({duplicate_rate:.1%})."
            )

        for col in df.columns:
            series = df[col]
            col_name = str(col)
            null_count = int(series.isna().sum())
            null_pct = (null_count / analyzed_rows) * 100.0
            n_unique = int(series.nunique(dropna=True))

            if null_pct > 50.0:
                anomalies.append(f"**Critical Missing Values**: Column `{col_name}` has {null_pct:.1f}% nulls.")
            elif null_pct > 20.0:
                anomalies.append(f"**Elevated Missing Values**: Column `{col_name}` has {null_pct:.1f}% nulls.")

            if n_unique == 1 and analyzed_rows > 1:
                anomalies.append(f"**Constant Column**: Column `{col_name}` has only 1 distinct value (zero variance).")

            # Logical Type Inference
            non_null_series = series.dropna()
            is_numeric = pd.api.types.is_numeric_dtype(series)

            if is_numeric and len(non_null_series) > 0:
                min_val = float(non_null_series.min())
                max_val = float(non_null_series.max())
                mean_val = float(non_null_series.mean())
                q25 = float(non_null_series.quantile(0.25))
                q75 = float(non_null_series.quantile(0.75))
                iqr = q75 - q25
                outliers_count = 0
                if iqr > 0:
                    lower_bound = q25 - 1.5 * iqr
                    upper_bound = q75 + 1.5 * iqr
                    outliers_count = int(((non_null_series < lower_bound) | (non_null_series > upper_bound)).sum())
                    if outliers_count > 0:
                        outlier_pct = (outliers_count / len(non_null_series)) * 100.0
                        if outlier_pct > 5.0:
                            anomalies.append(
                                f"**Significant Outliers**: Column `{col_name}` has {outliers_count} statistical outliers ({outlier_pct:.1f}%)."
                            )

                col_profiles.append({
                    "name": col_name,
                    "type": "Numeric",
                    "null_pct": f"{null_pct:.1f}%",
                    "distinct": n_unique,
                    "stats": f"min: {min_val:.2f}, max: {max_val:.2f}, mean: {mean_val:.2f}, outliers: {outliers_count}",
                })
            else:
                top_vals = non_null_series.value_counts().head(2).to_dict()
                top_str = ", ".join([f"'{k}': {v}" for k, v in top_vals.items()]) if top_vals else "N/A"
                col_profiles.append({
                    "name": col_name,
                    "type": "Categorical/Text",
                    "null_pct": f"{null_pct:.1f}%",
                    "distinct": n_unique,
                    "stats": f"top values: [{top_str}]",
                })

        # Calculate Overall Health Score
        completeness_score = completeness_rate * 50.0
        uniqueness_score = (1.0 - duplicate_rate) * 30.0
        anomaly_penalty = min(20.0, len(anomalies) * 4.0)
        validity_score = 20.0 - anomaly_penalty
        overall_score = round(completeness_score + uniqueness_score + validity_score, 1)

        if overall_score >= 90.0:
            verdict = "EXCELLENT (Production Ready)"
        elif overall_score >= 75.0:
            verdict = "GOOD (Minor Cleaning Needed)"
        elif overall_score >= 50.0:
            verdict = "NEEDS ATTENTION (Notable Quality Issues)"
        else:
            verdict = "CRITICAL (Severe Defects Detected)"

        # Build Markdown Report
        lines = [
            "# Data Quality & Governance Profile Report",
            f"**Overall Quality Score:** `{overall_score} / 100` — **{verdict}**",
            "",
            "## 1. Executive Summary",
            f"- **Total Rows:** {total_rows:,}" + (f" *(Profiled sample: {analyzed_rows:,})*" if is_sampled else ""),
            f"- **Total Columns:** {total_cols}",
            f"- **Completeness Rate:** `{completeness_rate:.1%}` ({missing_cells:,} missing cells)",
            f"- **Row Duplication Rate:** `{duplicate_rate:.1%}` ({duplicate_rows_count:,} duplicate rows)",
            f"- **Primary Key Candidates:** {', '.join([f'`{k}`' for k in pk_candidates]) if pk_candidates else 'None detected (No single unique non-null column)'}",
            "",
            "## 2. Key Actionable Anomalies",
        ]

        if anomalies:
            for anomaly in anomalies:
                lines.append(f"- {anomaly}")
        else:
            lines.append("- No critical quality anomalies detected! Data conforms to expected bounds.")

        lines.extend([
            "",
            "## 3. Column-by-Column Health Breakdown",
            "| Column Name | Inferred Type | Missing % | Unique Count | Distribution / Top Values |",
            "| :--- | :--- | :---: | :---: | :--- |",
        ])

        for cp in col_profiles:
            lines.append(f"| `{cp['name']}` | {cp['type']} | {cp['null_pct']} | {cp['distinct']:,} | {cp['stats']} |")

        lines.extend([
            "",
            "---",
            "*Report generated deterministically according to DAMA-DMBOK / CDMP® standards via `smolagents` Data Quality Profiler.*",
        ])

        return "\n".join(lines)

    def __init__(self, *args, **kwargs):
        self.is_initialized = False