main.py

"""
AAE -- Autonomous Analytics Engineer
=====================================
FastAPI server with file upload, WebSocket pipeline, and 5-phase agents.
"""

import asyncio
import json
import math
import os
import shutil
import re
from pathlib import Path

import numpy as np
import pandas as pd

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File, Body
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from agents import (
    AuditorAgent, CleanerAgent, AnalyzerAgent, ArchitectAgent, StorytellerAgent,
    AssuranceAgent,
    ChunkedAuditorAgent, ChunkedCleanerAgent, ChunkedAnalyzerAgent, ChunkedArchitectAgent,
)
from agents.auditor import _load_file, _detect_date_columns, _classify_column, _type_recommendation
from agents.io_utils import read_tabular_file, readiness_profile
BASE = Path(__file__).parent
DATA_DIR = BASE / "data"
OUTPUT_DIR = BASE / "output"
MAX_UPLOAD_BYTES = 5 * 1024 * 1024 * 1024
LARGE_FILE_BYTES = 250 * 1024 * 1024
PREVIEW_SAMPLE_ROWS = 20_000
CSV_CHUNK_ROWS = 50_000
SUPPORTED_UPLOADS = {".csv", ".xlsx", ".xls", ".parquet", ".pq", ".tsv"}
SAMPLEABLE_UPLOADS = {".csv", ".tsv"}
NUMERIC_TEXT_CLEAN_RE = re.compile(r"[,\$\u00a3\u20ac\u20b9%]")

app = FastAPI(title="Autonomous Analytics Engineer")
app.mount("/static", StaticFiles(directory=BASE / "static"), name="static")


# -- State --
pipeline_state = {
    "status": "idle",
    "current_phase": 0,
    "reports": {},
    "uploaded_file": None,
    "config": {},
}

REPORT_FILES = {
    "audit": "data_health_report.json",
    "clean": "cleaning_report.json",
    "analyze": "statistical_analysis_report.json",
    "architect": "architect_report.json",
    "story": "dashboard_stories.json",
    "assurance": "assurance_report.json",
}


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer): return int(obj)
        if isinstance(obj, np.floating): return float(obj)
        if isinstance(obj, np.bool_): return bool(obj)
        if isinstance(obj, np.ndarray): return obj.tolist()
        return super().default(obj)


def _safe_json(data):
    return json.loads(json.dumps(data, cls=NpEncoder, default=str))


def _save_report(phase: str, report: dict) -> None:
    filename = REPORT_FILES.get(phase)
    if not filename:
        return
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    with open(OUTPUT_DIR / filename, "w", encoding="utf-8") as f:
        json.dump(_safe_json(report), f, indent=2)


def _load_report_from_disk(phase: str) -> dict | None:
    filename = REPORT_FILES.get(phase)
    if not filename:
        return None
    path = OUTPUT_DIR / filename
    if not path.exists():
        return None
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _resolve_data_source() -> Path:
    uploaded = pipeline_state.get("uploaded_file")
    if uploaded:
        uploaded_path = DATA_DIR / uploaded
        if uploaded_path.exists():
            return uploaded_path
        raise FileNotFoundError(f"Uploaded dataset is no longer available: {uploaded}")
    raise FileNotFoundError("Upload a dataset before running the analytics pipeline.")


def _file_meta(path: Path) -> dict:
    size_bytes = path.stat().st_size if path.exists() else 0
    ext = path.suffix.lower()
    is_large = size_bytes > LARGE_FILE_BYTES and ext in SAMPLEABLE_UPLOADS
    return {
        "file_size_mb": round(size_bytes / 1024 / 1024, 2),
        "is_large_file": bool(is_large),
        "processing_mode": "chunked_full_file" if is_large else "full_dataset",
    }


def _read_preview_file(path: Path) -> tuple[pd.DataFrame, dict]:
    meta = _file_meta(path)
    ext = path.suffix.lower()
    if meta["is_large_file"] and ext in SAMPLEABLE_UPLOADS:
        df = read_tabular_file(str(path), nrows=PREVIEW_SAMPLE_ROWS)
        meta["preview_is_sampled"] = True
        meta["preview_rows_scanned"] = int(len(df))
        return df, meta
    df = _load_file(str(path))
    meta["preview_is_sampled"] = False
    meta["preview_rows_scanned"] = int(len(df))
    return df, meta


async def _prepare_pipeline_source(raw_path: Path, send) -> tuple[Path, dict]:
    meta = _file_meta(raw_path)
    readiness = readiness_profile(str(raw_path), nrows=PREVIEW_SAMPLE_ROWS)
    meta["readiness"] = readiness
    if readiness.get("status") == "blocked":
        detail = "; ".join(i.get("detail", i.get("title", "")) for i in readiness.get("issues", [])[:3])
        raise ValueError(f"Dataset is not ready for analysis: {detail}")
    if readiness.get("issues"):
        await send({
            "phase": "SETUP",
            "phase_number": 0,
            "action": "READINESS",
            "detail": f"Readiness {readiness.get('readiness_score')}/100: {len(readiness.get('issues', []))} issue(s) detected and disclosed.",
            "progress": None,
            "timestamp": "",
        })
    if not meta["is_large_file"]:
        meta["analysis_is_sampled"] = False
        return raw_path, meta

    await send({
        "phase": "SETUP",
        "phase_number": 0,
        "action": "CHUNKED_WAREHOUSE",
        "detail": (
            f"Large CSV warehouse mode: scanning and transforming the full "
            f"{meta['file_size_mb']} MB file in chunks."
        ),
        "progress": None,
        "timestamp": "",
    })
    meta.update({
        "analysis_is_sampled": False,
        "chunked_full_file": True,
    })
    return raw_path, meta


def _is_currency_like(name: str) -> bool:
    return any(k in name.lower() for k in (
        "amount", "revenue", "sales", "total", "price", "cost", "income",
        "profit", "value", "fare", "fee", "payment"
    ))


QUALITY_KPI_KEYWORDS = {"score": 3.5, "rating": 3.5, "satisfaction": 3.5, "nps": 4, "quality": 3}
COUNT_LIKE_KEYWORDS = {
    "count", "number", "num", "qty", "quantity", "volume", "views", "votes",
    "members", "followers", "subscribers", "population", "scored_by", "rated_by",
}


def _token_match(lower: str, kw: str) -> bool:
    return lower == kw or lower.endswith("_" + kw) or lower.startswith(kw + "_") or f"_{kw}_" in lower


def _preview_kpi_score(df: pd.DataFrame, name: str) -> tuple[float, str]:
    s = _coerce_numeric_preview(df[name]).dropna()
    lower = name.lower()
    semantic = 0.0
    reasons = []
    for kw in ("amount", "revenue", "sales", "total", "price", "cost", "income", "profit", "value", "fare", "fee", "payment", "spend"):
        if kw in lower:
            semantic += 4.0
            reasons.append("business value name")
            break
    for kw, weight in QUALITY_KPI_KEYWORDS.items():
        if _token_match(lower, kw):
            semantic += weight
            reasons.append("quality/rating KPI name")
            break
    for kw in COUNT_LIKE_KEYWORDS:
        if kw in lower:
            semantic -= 4.0
            reasons.append("count/popularity field")
            break
    magnitude = math.log10(abs(float(s.sum())) + 1) if len(s) else 0.0
    variability = min(float(s.std() / (abs(float(s.mean())) + 1e-9)), 5.0) if len(s) > 2 else 0.0
    score = semantic * 3 + magnitude + variability
    reason = ", ".join(dict.fromkeys(reasons)) or "numeric measure profile"
    return score, f"Auto Analyst score {score:.2f}: {reason}."


def _coerce_numeric_preview(series: pd.Series) -> pd.Series:
    if pd.api.types.is_numeric_dtype(series):
        return pd.to_numeric(series, errors="coerce")
    cleaned = (
        series.dropna().astype(str)
        .str.replace(NUMERIC_TEXT_CLEAN_RE, "", regex=True)
        .str.strip()
    )
    return pd.to_numeric(cleaned, errors="coerce")


def _is_numeric_candidate(df: pd.DataFrame, col_meta: dict) -> bool:
    name = col_meta["name"]
    if name not in df.columns:
        return False
    if pd.api.types.is_numeric_dtype(df[name]):
        return True
    if col_meta.get("recommendation") == "convert_to_numeric":
        return True
    sample = _coerce_numeric_preview(df[name].head(500))
    return bool(len(sample) and sample.notna().mean() >= 0.75)


def _suggest_value_column(df: pd.DataFrame, columns: list[dict]) -> dict | None:
    numeric = [
        c for c in columns
        if _is_numeric_candidate(df, c) and c["role"] not in ("id", "categorical_numeric")
    ]
    if not numeric:
        return None
    ranked = sorted(
        ((c, *_preview_kpi_score(df, c["name"])) for c in numeric if c["name"] in df.columns),
        key=lambda x: x[1],
        reverse=True,
    )
    if not ranked:
        return None
    winner, _, reason = ranked[0]
    return {"column": winner["name"], "reason": reason}


def _interview_questions(df: pd.DataFrame, columns: list[dict], suggested_value: dict | None,
                         suggested_date: dict | None) -> list[dict]:
    numeric_cols = [
        c for c in columns
        if _is_numeric_candidate(df, c) and c["role"] not in ("id", "categorical_numeric")
    ]
    date_cols = [c for c in columns if c["role"] == "date" or c.get("recommendation") == "parse_as_datetime"]
    risky_cols = [
        c for c in columns
        if c["role"] == "id" or c["null_pct"] > 10 or c["cardinality_band"] == "high"
    ][:8]
    value_options = [c["name"] for c in sorted(
        numeric_cols,
        key=lambda c: (
            0 if suggested_value and c["name"] == suggested_value.get("column") else 1,
            -abs(float(_coerce_numeric_preview(df[c["name"]]).sum())) if c["name"] in df else 0,
        ),
    )[:6]]
    date_options = [c["name"] for c in date_cols[:5]]

    questions = [
        {
            "id": "objective",
            "type": "text",
            "question": "What decision should this analysis help someone make?",
            "placeholder": "e.g. explain revenue growth, identify churn drivers, monitor data quality",
            "maps_to": "business_context.objective",
            "reason": "Senior analysts start by anchoring the dataset to a decision, not just a chart.",
        },
        {
            "id": "primary_value",
            "type": "choice",
            "question": "Which numeric field is the main KPI?",
            "options": value_options,
            "default": suggested_value.get("column") if suggested_value else (value_options[0] if value_options else ""),
            "maps_to": "primary_value",
            "reason": suggested_value.get("reason") if suggested_value else "No obvious KPI was found, so your choice matters.",
        },
    ]
    if date_options:
        questions.append({
            "id": "primary_date",
            "type": "choice",
            "question": "Which date should define trend and time intelligence?",
            "options": date_options,
            "default": suggested_date.get("column") if suggested_date else date_options[0],
            "maps_to": "primary_date",
            "reason": "Different date fields can produce different trend narratives.",
        })
    questions.append({
        "id": "cleaning_governance",
        "type": "multi",
        "question": "Which high-impact cleaning actions may AAE perform automatically?",
        "options": [
            {"value": "impute_nulls", "label": "Impute missing values", "default": True},
            {"value": "clip_outliers", "label": "Clip extreme outliers", "default": False},
            {"value": "delete_rows", "label": "Delete duplicate/invalid rows", "default": False},
            {"value": "repair_computed", "label": "Repair computed columns", "default": False},
        ],
        "maps_to": "governance.approvals",
        "reason": "These actions change data, so AAE treats them as explicit analyst approvals.",
    })
    if risky_cols:
        questions.append({
            "id": "protected_columns",
            "type": "multi",
            "question": "Which columns should AAE avoid changing?",
            "options": [{"value": c["name"], "label": c["name"], "default": c["role"] == "id"} for c in risky_cols],
            "maps_to": "governance.protected_columns",
            "reason": "IDs, high-cardinality fields, and sparse fields are common places to preserve source truth.",
        })
    return questions[:5]


def _apply_pipeline_config(audit_report: dict, config: dict) -> dict:
    config = config or {}
    role_overrides = config.get("role_overrides") or {}
    if role_overrides:
        for col, role in role_overrides.items():
            if col in audit_report.get("column_roles", {}):
                audit_report["column_roles"][col] = role
            for meta in audit_report.get("columns", []):
                if meta.get("name") == col:
                    meta["role"] = role
                    meta["user_override"] = True

    primary_date = config.get("primary_date")
    if primary_date:
        audit_report.setdefault("summary", {}).setdefault("date_columns", [])
        if primary_date not in audit_report["summary"]["date_columns"]:
            audit_report["summary"]["date_columns"].insert(0, primary_date)
        else:
            dates = audit_report["summary"]["date_columns"]
            audit_report["summary"]["date_columns"] = [primary_date] + [d for d in dates if d != primary_date]
        audit_report.setdefault("column_roles", {})[primary_date] = "date"
        for meta in audit_report.get("columns", []):
            if meta.get("name") == primary_date:
                meta["role"] = "date"
                meta["user_override"] = True

    primary_value = config.get("primary_value")
    if primary_value:
        audit_report.setdefault("column_roles", {})[primary_value] = "measure"
        for meta in audit_report.get("columns", []):
            if meta.get("name") == primary_value:
                meta["role"] = "measure"
                meta["user_override"] = True

    audit_report["user_config"] = config
    return audit_report


# -- Routes --

def _fresh_pipeline_state() -> None:
    """Reset all session state and remove any leftover upload + report artifacts.

    Called on every page load so that reloading the site always starts the user
    in a clean state with no preloaded dataset, no stale run, and no leftover
    artifacts on disk. Skipped while a pipeline is actively running so that a
    misclick on Refresh does not destroy mid-run state.
    """
    if pipeline_state.get("status") == "running":
        return
    prior_upload = pipeline_state.get("uploaded_file")
    pipeline_state["status"] = "idle"
    pipeline_state["current_phase"] = 0
    pipeline_state["reports"] = {}
    pipeline_state["uploaded_file"] = None
    pipeline_state["config"] = {}
    if prior_upload:
        try:
            (DATA_DIR / prior_upload).unlink(missing_ok=True)
        except OSError:
            pass
    if OUTPUT_DIR.exists():
        for child in OUTPUT_DIR.iterdir():
            try:
                if child.is_file():
                    child.unlink(missing_ok=True)
                elif child.is_dir():
                    shutil.rmtree(child, ignore_errors=True)
            except OSError:
                continue


@app.get("/", response_class=HTMLResponse)
async def serve_dashboard():
    _fresh_pipeline_state()
    return (BASE / "static" / "index.html").read_text(encoding="utf-8")


@app.get("/api/state")
async def get_state():
    return JSONResponse({
        "status": pipeline_state["status"],
        "current_phase": pipeline_state["current_phase"],
        "uploaded_file": pipeline_state.get("uploaded_file"),
        "config": pipeline_state.get("config", {}),
    })


@app.post("/api/upload")
async def upload_file(file: UploadFile = File(...)):
    os.makedirs(DATA_DIR, exist_ok=True)
    original_name = file.filename or "uploaded_dataset"
    ext = os.path.splitext(original_name)[1].lower()
    if ext not in SUPPORTED_UPLOADS:
        return JSONResponse({"error": f"Unsupported file type: {ext}"}, status_code=400)

    safe_stem = re.sub(r"[^A-Za-z0-9._-]+", "_", Path(original_name).stem).strip("._") or "dataset"
    dest = DATA_DIR / f"uploaded_{safe_stem}{ext}"
    size_bytes = 0
    with open(dest, "wb") as f:
        while True:
            chunk = await file.read(1024 * 1024)
            if not chunk:
                break
            size_bytes += len(chunk)
            if size_bytes > MAX_UPLOAD_BYTES:
                f.close()
                try:
                    dest.unlink(missing_ok=True)
                except Exception:
                    pass
                return JSONResponse({"error": "File is larger than the 5GB upload limit."}, status_code=413)
            f.write(chunk)

    pipeline_state["uploaded_file"] = dest.name
    pipeline_state["config"] = {}
    meta = _file_meta(dest)
    meta["readiness"] = readiness_profile(str(dest), nrows=PREVIEW_SAMPLE_ROWS)
    size_kb = round(size_bytes / 1024, 1)
    return JSONResponse({"filename": dest.name, "original_filename": original_name,
                         "size_kb": size_kb, "path": str(dest), **meta})


@app.get("/api/preview")
async def dataset_preview():
    try:
        if not pipeline_state.get("uploaded_file"):
            return JSONResponse({"error": "No uploaded dataset selected."}, status_code=400)
        raw_path = _resolve_data_source()
        df, preview_meta = _read_preview_file(raw_path)
        readiness = readiness_profile(str(raw_path), nrows=PREVIEW_SAMPLE_ROWS)
        sample_df = df.head(8)
        date_cols = _detect_date_columns(df)
        columns = []
        for col in df.columns:
            classification = _classify_column(col, df[col], len(df))
            role = "date" if col in date_cols else classification["role"]
            recommendation = _type_recommendation(col, df[col], date_cols)
            sample_values = [None if pd.isna(v) else str(v) for v in df[col].dropna().head(4).tolist()]
            nulls = int(df[col].isna().sum())
            columns.append({
                "name": col,
                "dtype": str(df[col].dtype),
                "role": role,
                "cardinality": classification["cardinality"],
                "cardinality_band": classification["cardinality_band"],
                "cardinality_ratio": classification["cardinality_ratio"],
                "nulls": nulls,
                "null_pct": round(nulls / max(len(df), 1) * 100, 2),
                "recommendation": recommendation,
                "sample_values": sample_values,
            })
        suggested_date = {"column": date_cols[0], "reason": "Best parseable date/time column."} if date_cols else None
        suggested_value = _suggest_value_column(df, columns)
        return JSONResponse(_safe_json({
            "filename": raw_path.name,
            "rows": int(len(df)),
            "columns_count": int(len(df.columns)),
            "memory_mb": round(df.memory_usage(deep=True).sum() / 1024 / 1024, 2),
            **preview_meta,
            "readiness": readiness,
            "columns": columns,
            "preview_rows": json.loads(sample_df.to_json(orient="records", date_format="iso")),
            "suggested_date": suggested_date,
            "suggested_value": suggested_value,
            "interview_questions": _interview_questions(df, columns, suggested_value, suggested_date),
        }))
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=400)


@app.post("/api/config")
async def save_config(config: dict | None = Body(default=None)):
    pipeline_state["config"] = config or {}
    return JSONResponse({"status": "saved", "config": pipeline_state["config"]})


@app.post("/api/reset")
async def reset_pipeline():
    pipeline_state["status"] = "idle"
    pipeline_state["current_phase"] = 0
    pipeline_state["reports"] = {}
    pipeline_state["config"] = {}
    # Clear output dir
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    os.makedirs(OUTPUT_DIR / "star_schema", exist_ok=True)
    return JSONResponse({"status": "reset"})


@app.get("/api/reports/{phase}")
async def get_report(phase: str):
    if phase in pipeline_state["reports"]:
        return JSONResponse(_safe_json(pipeline_state["reports"][phase]))
    persisted = _load_report_from_disk(phase)
    if persisted is not None:
        pipeline_state["reports"][phase] = persisted
        return JSONResponse(_safe_json(persisted))
    return JSONResponse({"error": "Report not found"}, status_code=404)


@app.get("/api/artifacts")
async def list_artifacts():
    files = []
    if OUTPUT_DIR.exists():
        for p in OUTPUT_DIR.rglob("*.parquet"):
            files.append({"name": p.name, "path": str(p.relative_to(BASE)),
                           "size_kb": round(p.stat().st_size / 1024, 1)})
        for p in OUTPUT_DIR.rglob("*.json"):
            files.append({"name": p.name, "path": str(p.relative_to(BASE)),
                           "size_kb": round(p.stat().st_size / 1024, 1)})
        for p in OUTPUT_DIR.rglob("*.csv"):
            files.append({"name": p.name, "path": str(p.relative_to(BASE)),
                           "size_kb": round(p.stat().st_size / 1024, 1)})
    return JSONResponse(files)


@app.get("/api/download/{filepath:path}")
async def download_artifact(filepath: str):
    full = BASE / filepath
    if full.exists():
        return FileResponse(full, filename=full.name)
    return JSONResponse({"error": "File not found"}, status_code=404)


# -- WebSocket Pipeline --

@app.websocket("/ws/pipeline")
async def pipeline_ws(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            msg = await ws.receive_text()
            cmd = json.loads(msg)
            if cmd.get("action") == "run":
                if "config" in cmd:
                    pipeline_state["config"] = cmd.get("config") or {}
                await _run_pipeline(ws)
            elif cmd.get("action") == "status":
                await ws.send_json({"type": "state", **_safe_json(pipeline_state)})
    except WebSocketDisconnect:
        pass


async def _run_pipeline(ws: WebSocket):
    pipeline_state["status"] = "running"
    pipeline_state["current_phase"] = 0
    pipeline_state["reports"] = {}

    # Clear previous output
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    os.makedirs(OUTPUT_DIR / "star_schema", exist_ok=True)

    async def send(entry: dict):
        await ws.send_json({"type": "log", **_safe_json(entry)})

    try:
        # -- Determine data source --
        source_path = _resolve_data_source()
        raw_path_obj, processing_meta = await _prepare_pipeline_source(source_path, send)
        raw_path = str(raw_path_obj)

        await send({"phase": "SETUP", "phase_number": 0, "action": "SOURCE",
                    "detail": f"Data source: {os.path.basename(raw_path)}",
                    "progress": None, "timestamp": ""})

        chunked_full_file = bool(processing_meta.get("chunked_full_file"))

        # -- Phase 1: AUDIT --
        pipeline_state["current_phase"] = 1
        await ws.send_json({"type": "phase_change", "phase": 1})
        auditor_cls = ChunkedAuditorAgent if chunked_full_file else AuditorAgent
        auditor = auditor_cls(raw_path, str(OUTPUT_DIR))
        audit_report = await auditor.run(callback=send)
        audit_report.setdefault("summary", {})["original_source_file"] = os.path.basename(str(source_path))
        audit_report.setdefault("summary", {})["processing_mode"] = processing_meta.get("processing_mode")
        audit_report.setdefault("summary", {})["analysis_is_sampled"] = processing_meta.get("analysis_is_sampled", False)
        audit_report["summary"]["chunked_full_file"] = chunked_full_file
        audit_report["summary"]["readiness"] = processing_meta.get("readiness")
        audit_report = _apply_pipeline_config(audit_report, pipeline_state.get("config", {}))
        pipeline_state["reports"]["audit"] = audit_report
        _save_report("audit", audit_report)

        # -- Phase 2: CLEAN --
        pipeline_state["current_phase"] = 2
        await ws.send_json({"type": "phase_change", "phase": 2})
        cleaner_cls = ChunkedCleanerAgent if chunked_full_file else CleanerAgent
        cleaner = cleaner_cls(raw_path, audit_report, str(OUTPUT_DIR),
                              pipeline_state.get("config", {}))
        clean_report = await cleaner.run(callback=send)
        pipeline_state["reports"]["clean"] = clean_report
        _save_report("clean", clean_report)

        # -- Phase 3: ANALYZE --
        pipeline_state["current_phase"] = 3
        await ws.send_json({"type": "phase_change", "phase": 3})
        parquet_path = str(OUTPUT_DIR / "cleaned_data.parquet")
        analyzer_cls = ChunkedAnalyzerAgent if chunked_full_file else AnalyzerAgent
        analyzer = analyzer_cls(parquet_path, audit_report, str(OUTPUT_DIR),
                                pipeline_state.get("config", {}))
        analysis_report = await analyzer.run(callback=send)
        pipeline_state["reports"]["analyze"] = analysis_report
        _save_report("analyze", analysis_report)

        # -- Phase 4: ARCHITECT --
        pipeline_state["current_phase"] = 4
        await ws.send_json({"type": "phase_change", "phase": 4})
        if chunked_full_file:
            architect = ChunkedArchitectAgent(parquet_path, audit_report, analysis_report, str(OUTPUT_DIR),
                                             pipeline_state.get("config", {}))
        else:
            architect = ArchitectAgent(parquet_path, audit_report, str(OUTPUT_DIR),
                                       pipeline_state.get("config", {}))
        architect_report = await architect.run(callback=send)
        pipeline_state["reports"]["architect"] = architect_report
        _save_report("architect", architect_report)

        # -- Phase 5: STORY --
        pipeline_state["current_phase"] = 5
        await ws.send_json({"type": "phase_change", "phase": 5})
        storyteller = StorytellerAgent(audit_report, clean_report,
                                       analysis_report, architect_report, str(OUTPUT_DIR))
        story_report = await storyteller.run(callback=send)
        pipeline_state["reports"]["story"] = story_report
        _save_report("story", story_report)

        assurance = AssuranceAgent(audit_report, clean_report, analysis_report,
                                   architect_report, story_report, str(OUTPUT_DIR))
        assurance_report = await assurance.run(callback=send)
        pipeline_state["reports"]["assurance"] = assurance_report
        _save_report("assurance", assurance_report)

        pipeline_state["status"] = "complete"
        await ws.send_json({"type": "complete", "message": "Pipeline complete"})

    except Exception as e:
        pipeline_state["status"] = "error"
        import traceback
        tb = traceback.format_exc()
        print(tb)
        await ws.send_json({"type": "error", "message": str(e)})


if __name__ == "__main__":
    import uvicorn
    os.makedirs(DATA_DIR, exist_ok=True)
    os.makedirs(OUTPUT_DIR / "star_schema", exist_ok=True)
    print("\n  +--------------------------------------------+")
    print("  |  AAE -- Autonomous Analytics Engineer     |")
    print("  |  Dashboard -> http://localhost:8000        |")
    print("  +--------------------------------------------+\n")
    uvicorn.run(app, host="0.0.0.0", port=8000)