from __future__ import annotations

import pandas as pd

from storai.analysis.root_cause import analyze_root_causes
from storai.insights.detector import detect_insights
from storai.metrics.engine import compute_metrics
from storai.ml.engine import detect_ml_findings
from storai.narrative.generator import generate_story
from storai.schema.infer import infer_schema


def run_analysis(summary_df, source_df=None, source_sql=None, db_engine=None):
    prepared_summary = prepare_summary_df(summary_df)

    schema = infer_schema(prepared_summary)
    metrics = compute_metrics(prepared_summary, schema)
    insights = detect_insights(metrics)

    resolved_source_df = _resolve_source_for_root_cause(
        source_df=source_df,
        source_sql=source_sql,
        db_engine=db_engine,
    )
    if resolved_source_df is not None:
        try:
            insights = analyze_root_causes(insights, resolved_source_df)
        except Exception as e:
            print(f"Root cause analysis skipped: {e}")

    insight_output = generate_story(insights)
    ml_findings = detect_ml_findings(prepared_summary, schema)

    return {
        "display_df": prepared_summary.copy(),
        "display_rows": prepared_summary.to_dict(orient="records"),
        "source_sql": source_sql or "",
        "insights": insight_output,
        "ml_findings": ml_findings,
    }


def prepare_summary_df(summary_df: pd.DataFrame) -> pd.DataFrame:
    prepared = summary_df.copy()

    # Build a canonical monthly date so growth/trend can use a true time axis.
    if "report_date" not in prepared.columns:
        if "year" in prepared.columns and "month" in prepared.columns:
            prepared["report_date"] = pd.to_datetime(
                {
                    "year": prepared["year"],
                    "month": prepared["month"],
                    "day": 1,
                },
                errors="coerce",
            )

    return prepared


def _resolve_source_for_root_cause(source_df=None, source_sql=None, db_engine=None):
    if source_df is not None:
        return source_df

    if source_sql and db_engine:
        return pd.read_sql(source_sql, db_engine)

    return None
