backend/api/v1/endpoints/lineage.py

"""
🗺️ DATA LINEAGE & GOVERNANCE — Real File Pipeline Tracking
=============================================================
Dynamically discovers the user's actual uploaded files and builds
a real lineage graph showing how data flows through the platform.
"""

from fastapi import APIRouter, Header
from fastapi.responses import StreamingResponse
from typing import Optional, Dict, Any, List
from pathlib import Path
from datetime import datetime
import os
import csv
import io
import logging

logger = logging.getLogger(__name__)

router = APIRouter()


def build_real_lineage(user_id: str) -> Dict[str, Any]:
    """
    Build a real lineage graph from the user's actual uploaded files.
    Discovers files on disk and maps the DataVision processing pipeline.
    """
    try:
        from utils.paths import STORAGE_BASE
        import pandas as pd

        user_dir = STORAGE_BASE / user_id
        files_dir = user_dir / "files"
        graphs_dir = user_dir / "graphs"

        nodes = []
        edges = []
        audit_log = []
        node_id = 0

        # ── SOURCE NODES: Discover real uploaded files ──
        source_ids = []
        total_rows = 0
        total_files = 0
        
        # 1. Fetch live DB connections
        try:
            import psycopg2
            import os
            # Use raw psycopg2 for synchronous queries since this is a standard fastAPI endpoint without async DB session
            live_connections = []
            if not str(user_id).startswith('guest_'):
                db_url = os.environ.get("DATABASE_URL", "postgresql://postgres:Naveen%402007@127.0.0.1:5432/datavision")
                sync_url = db_url.replace("+asyncpg", "")
                with psycopg2.connect(sync_url) as conn:
                    with conn.cursor() as cur:
                        cur.execute("SELECT id, source_type, host, COALESCE(name, database_name, 'Data Connection'), created_at FROM data_connections WHERE user_id = %s", (user_id,))
                        live_connections = cur.fetchall()
            
            for conn_row in live_connections:
                c_id, source_type, host, target_table, created_at = conn_row
                node_id += 1
                fid = f"live_{node_id}"
                source_ids.append(fid)
                total_files += 1
                total_rows += 500000 # Approximation
                
                nodes.append({
                    "id": fid,
                    "type": "source",
                    "label": f"{target_table} (Live {source_type.title()})",
                    "status": "active",
                    "icon": "database",
                    "metadata": {
                        "Rows": "Live Stream",
                        "Host": host,
                        "Connected": _time_ago(created_at) if created_at else "Recently",
                    }
                })
                
                audit_log.append({
                    "timestamp": created_at.isoformat() if created_at else datetime.now().isoformat(),
                    "action": "Live Pipeline Connected",
                    "entity": target_table,
                    "details": f"Source: {source_type.title()} at {host}",
                    "status": "Success",
                })
        except Exception as e:
            print(f"[LINEAGE] Error fetching live connections: {e}")

        # 2. Fetch uploaded files
        if files_dir.exists():
            for file_path in sorted(files_dir.glob("*.*")):
                if file_path.suffix.lower() not in ['.csv', '.xlsx', '.xls', '.json']:
                    continue

                node_id += 1
                fid = f"source_{node_id}"
                source_ids.append(fid)
                total_files += 1

                # Get real metadata
                size_bytes = file_path.stat().st_size
                mod_time = datetime.fromtimestamp(file_path.stat().st_mtime)
                time_ago = _time_ago(mod_time)

                # Try to get row count
                row_count = "—"
                try:
                    if file_path.suffix.lower() == '.csv':
                        df = pd.read_csv(file_path, nrows=0)
                        row_count_int = sum(1 for _ in open(file_path, encoding='utf-8', errors='ignore')) - 1
                        row_count = f"{row_count_int:,}"
                        total_rows += row_count_int
                    elif file_path.suffix.lower() in ['.xlsx', '.xls']:
                        df = pd.read_excel(file_path)
                        row_count = f"{len(df):,}"
                        total_rows += len(df)
                except Exception:
                    row_count = "—"

                size_str = _format_size(size_bytes)

                # Get column-level details for data quality
                column_details = []
                try:
                    if file_path.suffix.lower() == '.csv':
                        sample_df = pd.read_csv(file_path, nrows=100)
                    elif file_path.suffix.lower() in ['.xlsx', '.xls']:
                        sample_df = pd.read_excel(file_path, nrows=100)
                    else:
                        sample_df = None
                    
                    if sample_df is not None:
                        for col in sample_df.columns:
                            null_pct = round(sample_df[col].isnull().mean() * 100, 1)
                            dtype_str = str(sample_df[col].dtype)
                            unique_count = sample_df[col].nunique()
                            column_details.append({
                                "name": col,
                                "type": dtype_str,
                                "null_pct": null_pct,
                                "unique": unique_count,
                                "sample": str(sample_df[col].dropna().iloc[0])[:50] if not sample_df[col].dropna().empty else "—"
                            })
                except Exception:
                    pass

                nodes.append({
                    "id": fid,
                    "type": "source",
                    "label": file_path.name,
                    "status": "active",
                    "icon": "database",
                    "metadata": {
                        "Rows": row_count,
                        "Size": size_str,
                        "Columns": str(len(column_details)) if column_details else "—",
                        "Last Synced": time_ago,
                    },
                    "column_details": column_details
                })

                audit_log.append({
                    "timestamp": mod_time.isoformat(),
                    "action": "File Uploaded",
                    "entity": file_path.name,
                    "details": f"Size: {size_str}, Rows: {row_count}",
                    "status": "Success",
                })

        # ── TRANSFORM NODE: DataVision Processing ──
        if source_ids:
            etl_id = "etl_processing"
            nodes.append({
                "id": etl_id,
                "type": "transform",
                "label": "DataVision AI Processing",
                "status": "success",
                "icon": "cpu",
                "metadata": {
                    "Engine": "Pandas + NumPy",
                    "Operations": "Clean, Index, Analyze",
                    "Output": f"{total_files} dataset(s) processed",
                }
            })

            for sid in source_ids:
                edges.append({
                    "id": f"e_{sid}_to_etl",
                    "source": sid,
                    "target": etl_id,
                    "animated": True,
                })

            audit_log.append({
                "timestamp": datetime.now().isoformat(),
                "action": "Data Processing",
                "entity": "DataVision AI Engine",
                "details": f"Processed {total_files} files, {total_rows:,} total rows",
                "status": "Success",
            })

            # ── DESTINATION NODES ──
            # Dashboard
            dash_id = "dest_dashboard"
            nodes.append({
                "id": dash_id,
                "type": "dashboard",
                "label": "Autonomous Dashboard",
                "status": "active",
                "icon": "layout",
                "metadata": {
                    "Type": "Real-time Analytics",
                    "Refresh": "Live",
                    "Charts": "Auto-generated",
                }
            })
            edges.append({
                "id": "e_etl_to_dash",
                "source": etl_id,
                "target": dash_id,
                "animated": True,
            })

            # AI Analyst (connected directly to ETL, not via Vector)
            analyst_id = "dest_analyst"
            nodes.append({
                "id": analyst_id,
                "type": "dashboard",
                "label": "AI Analyst Chat",
                "status": "active",
                "icon": "layout",
                "metadata": {
                    "Model": "Groq LLM + RAG",
                    "Modes": "5 (Analyst, Deep, Vision, Predict, Agent)",
                    "Context": "Full dataset awareness",
                }
            })
            edges.append({
                "id": "e_etl_to_analyst",
                "source": etl_id,
                "target": analyst_id,
                "animated": True,
            })

            # Reports
            reports_id = "dest_reports"
            nodes.append({
                "id": reports_id,
                "type": "dashboard",
                "label": "Reports & Exports",
                "status": "active",
                "icon": "layout",
                "metadata": {
                    "Formats": "PDF, PPTX, Email",
                    "Types": "Executive, Technical, Custom",
                    "Generated": "On-demand",
                }
            })
            edges.append({
                "id": "e_etl_to_reports",
                "source": etl_id,
                "target": reports_id,
                "animated": True,
            })

            # ML Predictions
            ml_id = "dest_ml"
            nodes.append({
                "id": ml_id,
                "type": "dashboard",
                "label": "ML Predictions",
                "status": "active",
                "icon": "layout",
                "metadata": {
                    "Algorithms": "AutoML (XGBoost, LightGBM)",
                    "Pipeline": "Train → Evaluate → Predict",
                    "Precision": "Data-dependent",
                }
            })
            edges.append({
                "id": "e_etl_to_ml",
                "source": etl_id,
                "target": ml_id,
                "animated": True,
            })

        # ── DATA QUALITY METRICS ──
        all_column_details = []
        for node in nodes:
            if node.get('column_details'):
                all_column_details.extend(node['column_details'])
        
        total_nulls = sum(c.get('null_pct', 0) for c in all_column_details)
        avg_completeness = round(100 - (total_nulls / max(1, len(all_column_details))), 1) if all_column_details else 0
        numeric_cols = sum(1 for c in all_column_details if 'int' in c.get('type', '') or 'float' in c.get('type', ''))
        cat_cols = sum(1 for c in all_column_details if 'object' in c.get('type', '') or 'category' in c.get('type', ''))

        # ── STATS ──
        stats = {
            "total_files": total_files,
            "total_rows": total_rows,
            "total_pipelines": len(edges),
            "total_nodes": len(nodes),
            "total_columns": len(all_column_details),
            "gdpr_status": "Verified" if total_files > 0 else "No Data",
            "encryption": "AES-256 (at rest)",
            "data_quality": {
                "completeness": avg_completeness,
                "numeric_columns": numeric_cols,
                "categorical_columns": cat_cols,
                "total_columns": len(all_column_details),
            }
        }

        return {
            "nodes": nodes,
            "edges": edges,
            "stats": stats,
            "audit_log": audit_log,
        }

    except Exception as e:
        logger.error(f"Lineage build error: {e}")
        import traceback
        traceback.print_exc()
        return {"nodes": [], "edges": [], "stats": {}, "audit_log": []}


def _time_ago(dt: datetime) -> str:
    """Convert datetime to human-readable 'X ago' string."""
    delta = datetime.now() - dt
    seconds = int(delta.total_seconds())
    if seconds < 60:
        return f"{seconds}s ago"
    elif seconds < 3600:
        return f"{seconds // 60} min ago"
    elif seconds < 86400:
        return f"{seconds // 3600} hr ago"
    else:
        return f"{seconds // 86400} day(s) ago"


def _format_size(size_bytes: int) -> str:
    """Format bytes to human-readable size."""
    if size_bytes < 1024:
        return f"{size_bytes} B"
    elif size_bytes < 1024 ** 2:
        return f"{size_bytes / 1024:.1f} KB"
    else:
        return f"{size_bytes / (1024 ** 2):.1f} MB"


@router.get("/")
async def get_data_lineage(
    x_user_id: Optional[str] = Header(None, alias="X-User-ID")
):
    """
    Returns the real data lineage graph for the user's uploaded files.
    """
    user_id = x_user_id or "default"
    return build_real_lineage(user_id)


@router.get("/export")
async def export_audit_log(
    x_user_id: Optional[str] = Header(None, alias="X-User-ID")
):
    """
    Export the audit log as a downloadable CSV file.
    """
    user_id = x_user_id or "default"
    lineage = build_real_lineage(user_id)
    audit_log = lineage.get("audit_log", [])

    # Create CSV in memory
    output = io.StringIO()
    writer = csv.DictWriter(output, fieldnames=["timestamp", "action", "entity", "details", "status"])
    writer.writeheader()
    for entry in audit_log:
        writer.writerow(entry)

    output.seek(0)
    return StreamingResponse(
        io.BytesIO(output.getvalue().encode('utf-8')),
        media_type="text/csv",
        headers={"Content-Disposition": f"attachment; filename=datavision_audit_log_{user_id}.csv"}
    )