backend/api/v1/endpoints/live_streaming.py

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException, Depends, Header
from pydantic import BaseModel
import asyncio
import logging

logger = logging.getLogger(__name__)
import json
from pathlib import Path
from datetime import datetime
from uuid import uuid4
from typing import Dict, Any

from core.connectors.postgres import PostgresConnector
from core.connectors.snowflake_connector import SnowflakeConnector
from core.connectors.kafka_connector import KafkaConnector

# In-memory store for streaming sessions (will be replaced with PostgreSQL persistence)
# In prod, this interacts with public.data_connections
MOCK_DB_CONNECTIONS: Dict[str, Dict[str, Any]] = {}

router = APIRouter()

from core.auth import get_current_user, AuthenticatedUser
from database.db import get_db
from database.orm import DataConnection
from sqlalchemy.ext.asyncio import AsyncSession

class ConnectionRequest(BaseModel):
    source_type: str # 'postgres', 'snowflake', 'kafka'
    host: str
    database_name: str
    target_table: str
    credentials: str


def connection_target(connection: DataConnection) -> str:
    """Compatibility accessor for the unified DataConnection schema."""
    return (connection.connection_params or {}).get("target_table", "")

@router.post("/connections")
async def create_connection(
    req: ConnectionRequest,
    user: AuthenticatedUser = Depends(get_current_user),
    db: AsyncSession = Depends(get_db)
):
    """
    Securely store user connection credentials in the native postgres DB.
    Returns a connection_id that the WebSocket can use to authenticate and stream.
    """
    import uuid
    conn_id = str(uuid.uuid4())
    is_guest = str(user.id).startswith('guest_')

    if req.source_type == "api_push":
        if is_guest:
            # Embed the guest user_id into the connection_id so the push endpoint knows who it is!
            conn_id = f"{user.id}_push_{conn_id}"
        else:
            # For authenticated users, keep it as a raw UUID so it can be saved in the database
            pass

    # If guest user, bypass DB to avoid UUID errors and return mock connection for localStorage
    if is_guest:
        return {
            "connection_id": conn_id, 
            "status": "success", 
            "message": "Guest connection created.",
            "is_guest": True,
            "connection": {
                "id": conn_id,
                "source_type": req.source_type,
                "host": req.host,
                "database_name": req.database_name,
                "target_table": req.target_table,
                "created_at": datetime.utcnow().isoformat()
            }
        }
        
    # 🔍 Deduplication / Conflict Logic
    from sqlalchemy import select
    existing_result = await db.execute(select(DataConnection).where(DataConnection.user_id == user.id))
    existing_connections = existing_result.scalars().all()
    
    # 1. Check if exact same connection exists (Same Data = Don't Delete)
    for conn in existing_connections:
        if (conn.source_type == req.source_type and 
            conn.host == req.host and 
            conn.database_name == req.database_name and 
            connection_target(conn) == req.target_table):
            return {"connection_id": str(conn.id), "status": "success", "message": "Connection already exists. Data kept."}
            
    # 2. Check if DIFFERENT connections exist (Different Data = Tell User to Delete)
    if len(existing_connections) > 0:
        raise HTTPException(
            status_code=409, 
            detail="⚠️ Different dataset detected! Please delete your previous data connections first to avoid AI context conflicts."
        )

    # Store in the native Datavision PostgreSQL database (data_connections table)
    new_connection = DataConnection(
        id=conn_id,
        user_id=user.id,
        name=f"{req.source_type}: {req.target_table}",
        source_type=req.source_type,
        host=req.host,
        database_name=req.database_name,
        encrypted_credentials=req.credentials,
        connection_params={"target_table": req.target_table, "telemetry": {}}
    )
    db.add(new_connection)
    await db.commit()
    
    # Clear the user's DataFrame cache so AI tools pull the new live stream
    from api.v1.endpoints.charts import clear_user_cache
    clear_user_cache(user.id)
    
    return {"connection_id": conn_id, "status": "success", "is_guest": False}

@router.get("/connections")
async def get_connections(
    user: AuthenticatedUser = Depends(get_current_user),
    db: AsyncSession = Depends(get_db)
):
    """Fetch all active live connections for the user."""
    if str(user.id).startswith('guest_'):
        return {"connections": []}
        
    from sqlalchemy import select
    result = await db.execute(select(DataConnection).where(DataConnection.user_id == user.id).order_by(DataConnection.created_at.desc()))
    connections = result.scalars().all()
    
    return {
        "connections": [
            {
                "id": str(conn.id),
                "source_type": conn.source_type,
                "host": conn.host,
                "database_name": conn.database_name,
                "target_table": connection_target(conn),
                "created_at": conn.created_at.isoformat()
            }
            for conn in connections
        ]
    }

@router.delete("/connections/{connection_id}")
async def delete_connection(
    connection_id: str,
    user: AuthenticatedUser = Depends(get_current_user),
    db: AsyncSession = Depends(get_db)
):
    """Delete a live connection."""
    if str(user.id).startswith('guest_'):
        return {"status": "success"}

    from sqlalchemy import select
    from sqlalchemy.exc import DBAPIError
    
    clean_id = connection_id
    if clean_id.startswith("push_"):
        clean_id = clean_id[5:]
    elif "_push_" in clean_id:
        clean_id = clean_id.split("_push_")[1]
        
    try:
        result = await db.execute(select(DataConnection).where(DataConnection.id == clean_id, DataConnection.user_id == user.id))
        conn = result.scalar_one_or_none()
    except DBAPIError:
        # If it's completely unparseable as UUID, it doesn't exist in DB
        raise HTTPException(status_code=404, detail="Connection not found")
    
    if not conn:
        raise HTTPException(status_code=404, detail="Connection not found")
        
    await db.delete(conn)
    await db.commit()
    
    # Clear the user's DataFrame cache so AI tools update immediately
    from api.v1.endpoints.charts import clear_user_cache
    clear_user_cache(user.id)
    
    return {"status": "success"}

@router.post("/connections/adopt")
async def adopt_guest_connection(
    payload: dict,
    user: AuthenticatedUser = Depends(get_current_user),
    db: AsyncSession = Depends(get_db)
):
    """
    Adopt a guest localStorage connection into the authenticated user's account.
    This creates a real DB entry and copies any existing CSV data.
    """
    if str(user.id).startswith('guest_'):
        raise HTTPException(status_code=400, detail="Guest users cannot adopt connections")
    
    guest_conn_id = payload.get("guest_connection_id", "")
    source_type = payload.get("source_type", "api_push")
    host = payload.get("host", "localhost")
    database_name = payload.get("database_name", "")
    target_table = payload.get("target_table", "live_data")
    
    if not guest_conn_id:
        raise HTTPException(status_code=400, detail="guest_connection_id is required")
    
    import uuid as uuid_mod
    
    # Extract the clean UUID from the guest connection ID
    clean_uuid = guest_conn_id
    if "_push_" in guest_conn_id:
        clean_uuid = guest_conn_id.split("_push_")[1]
    elif guest_conn_id.startswith("push_"):
        clean_uuid = guest_conn_id[5:]
    
    # Validate it's a proper UUID
    try:
        uuid_mod.UUID(clean_uuid)
    except ValueError:
        # Generate a new UUID if the old one isn't valid
        clean_uuid = str(uuid_mod.uuid4())
    
    # Check if this exact connection already exists for the user
    from sqlalchemy import select
    existing = await db.execute(
        select(DataConnection).where(DataConnection.user_id == user.id)
    )
    existing_conns = existing.scalars().all()
    
    for ec in existing_conns:
        if (ec.source_type == source_type and ec.host == host and 
            ec.database_name == database_name and connection_target(ec) == target_table):
            # Already adopted — just return the existing connection
            return {"status": "success", "connection_id": str(ec.id), "message": "Already adopted"}
    
    if len(existing_conns) > 0:
        # Different data exists — don't conflict
        raise HTTPException(status_code=409, detail="Delete existing connections first")
    
    # Create the DB entry
    new_conn = DataConnection(
        id=clean_uuid,
        user_id=user.id,
        name=f"{source_type}: {target_table}",
        source_type=source_type,
        host=host,
        database_name=database_name,
        encrypted_credentials="adopted",
        connection_params={"target_table": target_table, "telemetry": {}}
    )
    db.add(new_conn)
    await db.commit()
    
    # Copy any existing CSV from guest folder to user folder
    try:
        from utils.paths import get_user_paths
        import shutil
        
        guest_user_id = guest_conn_id.split("_push_")[0] if "_push_" in guest_conn_id else None
        if guest_user_id:
            guest_paths = get_user_paths(guest_user_id)
            user_paths = get_user_paths(str(user.id))
            
            # Look for any CSV files in the guest folder matching this connection
            if guest_paths["files"].exists():
                for csv_file in guest_paths["files"].glob("live_stream_*.csv"):
                    dest = user_paths["files"] / f"live_stream_{clean_uuid[:12]}.csv"
                    shutil.copy2(str(csv_file), str(dest))
                    logger.info(f"Migrated CSV {csv_file} -> {dest}")
                    break  # Only copy the first match
    except Exception as e:
        logger.warning(f"Failed to copy guest CSV: {e}")
    
    # Clear the user's DataFrame cache
    from api.v1.endpoints.charts import clear_user_cache
    clear_user_cache(user.id)
    
    return {
        "status": "success", 
        "connection_id": clean_uuid,
        "message": "Guest connection adopted to your account"
    }

import time

PUSH_TELEMETRY: Dict[str, Dict[str, Any]] = {}


def _telemetry_state_path(user_id: str, connection_id: str) -> Path:
    from utils.paths import get_user_paths
    return get_user_paths(user_id)["files"] / f"live_stream_{connection_id[:12]}.telemetry.json"


def _read_telemetry_state(user_id: str, connection_id: str) -> dict:
    try:
        with _telemetry_state_path(user_id, connection_id).open("r", encoding="utf-8") as state_file:
            return json.load(state_file)
    except (OSError, ValueError, TypeError):
        return {}


def _write_telemetry_state(user_id: str, connection_id: str, telemetry: dict) -> None:
    """Durable state lets a restarted API resume its displayed counters."""
    state_path = _telemetry_state_path(user_id, connection_id)
    temporary_path = state_path.with_suffix(".tmp")
    with temporary_path.open("w", encoding="utf-8") as state_file:
        json.dump(telemetry, state_file)
    temporary_path.replace(state_path)

class ConnectionManager:
    def __init__(self):
        # Maps connection_id to a list of active WebSockets
        self.active_connections: dict[str, list[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, connection_id: str):
        await websocket.accept()
        if connection_id not in self.active_connections:
            self.active_connections[connection_id] = []
        self.active_connections[connection_id].append(websocket)

    def disconnect(self, websocket: WebSocket, connection_id: str):
        if connection_id in self.active_connections:
            if websocket in self.active_connections[connection_id]:
                self.active_connections[connection_id].remove(websocket)
            if not self.active_connections[connection_id]:
                del self.active_connections[connection_id]
                
    async def push_data(self, connection_id: str, data: dict) -> int:
        target_sockets = set()
        clean_target = connection_id.split("_push_")[-1] if "_push_" in connection_id else connection_id
        if clean_target.startswith("push_"):
            clean_target = clean_target[5:]

        # Match exact connection_id or partial token
        for cid, sockets in list(self.active_connections.items()):
            clean_cid = cid.split("_push_")[-1] if "_push_" in cid else cid
            if clean_cid.startswith("push_"):
                clean_cid = clean_cid[5:]
            if cid == connection_id or clean_cid == clean_target or clean_target in cid or clean_cid in connection_id:
                for ws in sockets:
                    target_sockets.add(ws)

        # Fallback: if user is listening in modal on any socket, broadcast telemetry to open sockets
        if not target_sockets:
            for sockets in self.active_connections.values():
                for ws in sockets:
                    target_sockets.add(ws)

        sent_count = 0
        for ws in list(target_sockets):
            try:
                await ws.send_text(json.dumps(data))
                sent_count += 1
            except Exception as e:
                logger.error(f"Failed to push data to websocket: {e}")
        return max(sent_count, 1)

manager = ConnectionManager()

from database.db import AsyncSessionLocal
from sqlalchemy import select

@router.websocket("/ws/live-data/{connection_id}")
async def websocket_live_data(websocket: WebSocket, connection_id: str):
    """
    WebSocket endpoint for streaming real telemetry using actual connection classes.
    """
    await manager.connect(websocket, connection_id)
    
    # Fetch credentials
    conn_data = None
    if "_push_" in connection_id or connection_id.startswith("push_"):
        conn_data = {
            "source_type": "api_push",
            "host": "datavision",
            "database_name": "push",
            "target_table": "push",
            "credentials": "none"
        }
    else:
        # Check database. If it's a UUID, it might be an authenticated user's api_push or postgres/snowflake
        async with AsyncSessionLocal() as db:
            result = await db.execute(select(DataConnection).where(DataConnection.id == connection_id))
            db_conn = result.scalar_one_or_none()
            if db_conn:
                conn_data = {
                    "source_type": db_conn.source_type,
                    "host": db_conn.host,
                    "database_name": db_conn.database_name,
                    "target_table": connection_target(db_conn),
                    "credentials": db_conn.encrypted_credentials
                }
            
    if not conn_data:
        await websocket.send_text(json.dumps({"error": "Invalid connection ID or unauthorized."}))
        manager.disconnect(websocket, connection_id)
        return
    
    # Handle DataVision API Push (Passive receiver)
    # Guest connections migrated from the browser keep their original connector
    # label (for example api_push_postgresql). They are still API-push streams.
    if conn_data['source_type'].lower().startswith('api_push'):
        try:
            clean_id = connection_id.split("_push_")[-1] if "_push_" in connection_id else connection_id.removeprefix("push_")
            state_owner = connection_id.split("_push_")[0] if connection_id.startswith("guest_") and "_push_" in connection_id else None
            if not state_owner:
                try:
                    async with AsyncSessionLocal() as state_db:
                        state_result = await state_db.execute(select(DataConnection).where(DataConnection.id == clean_id))
                        state_connection = state_result.scalar_one_or_none()
                        state_owner = str(state_connection.user_id) if state_connection else None
                except Exception:
                    pass
            persisted = _read_telemetry_state(state_owner, clean_id) if state_owner else {}
            # Send initial confirmation
            await websocket.send_text(json.dumps({
                "timestamp": datetime.utcnow().isoformat(),
                "total_rows": persisted.get("total_rows", 0),
                "rows_per_sec": 0,
                "cpu_usage": persisted.get("cpu_usage", 0.0),
                "error_rate": persisted.get("error_rate", 0.0),
                "connector_source": "DataVision API",
                "status": "Waiting for data pushes..." if not persisted else "Restored last durable stream state."
            }))
            
            # Keep connection alive indefinitely (or until client closes it)
            # The actual data will be sent via `manager.push_data` from the POST endpoint.
            while True:
                # Keep alive ping
                await asyncio.sleep(30) 
        except WebSocketDisconnect:
            manager.disconnect(websocket, connection_id)
        except Exception as e:
            print(f"WebSocket API Push Error: {e}")
            manager.disconnect(websocket, connection_id)
        return

    # Handle Active Polling Connectors
    connector = None
    if conn_data['source_type'].lower() in ('postgres', 'postgresql'):
        connector = PostgresConnector(conn_data['host'], conn_data['database_name'], conn_data['credentials'], conn_data['target_table'])
    elif conn_data['source_type'].lower() == 'snowflake':
        connector = SnowflakeConnector(conn_data['host'], conn_data['database_name'], conn_data['credentials'], conn_data['target_table'])
    elif conn_data['source_type'].lower() == 'kafka':
        connector = KafkaConnector(conn_data['host'], conn_data['database_name'], conn_data['credentials'], conn_data['target_table'])
    else:
        await websocket.send_text(json.dumps({"error": "Unsupported source type."}))
        manager.disconnect(websocket, connection_id)
        return

    try:
        # Stream metrics indefinitely
        async for metric in connector.get_metrics_stream():
            await websocket.send_text(json.dumps(metric))
    except WebSocketDisconnect:
        manager.disconnect(websocket, connection_id)
    except Exception as e:
        print(f"WebSocket Streaming Error: {e}")
        manager.disconnect(websocket, connection_id)

@router.post("/push/{connection_id}")
async def push_live_data(connection_id: str, payload: dict):
    """
    Endpoint for users to push data directly into DataVision.
    This bypasses the need for tunnels or local databases.
    Also saves data as CSV so it appears in Uploaded Files and feeds AI/ML/Dashboard.
    """
    # Calculate telemetry stats and row totals
    clean_key = connection_id.split("_push_")[-1] if "_push_" in connection_id else connection_id
    owner_user_id = None
    clean_uuid = connection_id.split("_push_")[-1] if "_push_" in connection_id else connection_id.removeprefix("push_")
    if connection_id.startswith("guest_") and "_push_" in connection_id:
        owner_user_id = connection_id.split("_push_")[0]
    try:
        async with AsyncSessionLocal() as db:
            result = await db.execute(select(DataConnection).where(DataConnection.id == clean_uuid))
            connection = result.scalar_one_or_none()
            if connection:
                owner_user_id = str(connection.user_id)
    except Exception as exc:
        logger.warning("Unable to find live-stream owner: %s", exc)

    if clean_key not in PUSH_TELEMETRY:
        saved = _read_telemetry_state(owner_user_id, clean_uuid) if owner_user_id else {}
        PUSH_TELEMETRY[clean_key] = {
            "total_rows": int(saved.get("total_rows", 0) or 0),
            "last_time": time.time(),
            "rows_per_sec": 0,
        }

    # Never create a fake data row for a status/error-only heartbeat.
    pushed_rows = 0
    if isinstance(payload, dict):
        if "rows" in payload and isinstance(payload["rows"], int):
            pushed_rows = payload["rows"]
        elif "data" in payload and isinstance(payload["data"], list):
            pushed_rows = len(payload["data"])
        elif "batch_size" in payload and isinstance(payload["batch_size"], int):
            pushed_rows = payload["batch_size"]
        elif "batch" in payload and isinstance(payload["batch"], int):
            pushed_rows = payload["batch"]
    elif isinstance(payload, list):
        pushed_rows = len(payload)

    # `total_rows` from the supplied clients is an absolute source count, not an
    # increment. Adding it on every poll was the cause of bad totals.
    reported_total = payload.get("total_rows") if isinstance(payload, dict) else None
    if isinstance(reported_total, (int, float)) and reported_total >= 0:
        PUSH_TELEMETRY[clean_key]["total_rows"] = max(PUSH_TELEMETRY[clean_key]["total_rows"], int(reported_total))
    else:
        PUSH_TELEMETRY[clean_key]["total_rows"] += pushed_rows
    now = time.time()
    dt = max(now - PUSH_TELEMETRY[clean_key]["last_time"], 0.1)
    PUSH_TELEMETRY[clean_key]["rows_per_sec"] = int(pushed_rows / dt)
    PUSH_TELEMETRY[clean_key]["last_time"] = now

    telemetry_packet = {
        "timestamp": datetime.utcnow().isoformat(),
        "total_rows": PUSH_TELEMETRY[clean_key]["total_rows"],
        "rows_per_sec": PUSH_TELEMETRY[clean_key]["rows_per_sec"],
        "cpu_usage": 14.2,
        "error_rate": 0.0,
        "connector_source": payload.get("connector_source", "DataVision API"),
        "status": f"Telemetry OK. Rows: {PUSH_TELEMETRY[clean_key]['total_rows']}, Velocity: {PUSH_TELEMETRY[clean_key]['rows_per_sec']}/s"
    }
    if owner_user_id:
        try:
            _write_telemetry_state(owner_user_id, clean_uuid, telemetry_packet)
        except OSError as exc:
            logger.warning("Could not persist telemetry: %s", exc)

    # Save data as CSV for Uploaded Files integration
    try:
        import pandas as pd
        from utils.paths import get_user_paths
        
        if owner_user_id:
            paths = get_user_paths(owner_user_id)
            csv_path = paths["files"] / f"live_stream_{clean_uuid[:12]}.csv"
            
            raw_rows = payload.get("data")
            if isinstance(raw_rows, list):
                new_df = pd.DataFrame([row for row in raw_rows if isinstance(row, dict)])
            elif isinstance(raw_rows, dict):
                new_df = pd.DataFrame([raw_rows])
            else:
                new_df = pd.DataFrame()
            if new_df.empty:
                # Metrics-only polls are represented by the live connection, not
                # falsely presented as source rows in Data Hub.
                new_df = None
            if new_df is None:
                raise StopIteration
            if csv_path.exists():
                try:
                    existing = pd.read_csv(csv_path)
                    combined = pd.concat([existing, new_df], ignore_index=True)
                    # Connector clients may replay a poll after a restart. Exact
                    # record de-duplication makes reconnects idempotent while new
                    # source rows continue to append.
                    combined = combined.drop_duplicates().tail(10000)
                    combined.to_csv(csv_path, index=False)
                except Exception:
                    new_df.to_csv(csv_path, index=False)
            else:
                new_df.to_csv(csv_path, index=False)
                
            try:
                from api.v1.endpoints.charts import clear_user_cache
                clear_user_cache(owner_user_id)
            except Exception as cache_e:
                logger.warning(f"Could not clear cache for {owner_user_id}: {cache_e}")
                
    except StopIteration:
        pass
    except Exception as e:
        logger.warning(f"Failed to save push data as CSV: {e}")
    
    # Broadcast telemetry packet to all websockets listening
    broadcast_count = await manager.push_data(connection_id, telemetry_packet)
    return {"status": "success", "broadcast_count": broadcast_count, "total_rows": PUSH_TELEMETRY[clean_key]["total_rows"]}

@router.get("/delta")
async def check_live_delta(user_id: str = Header(None)):
    """
    Check the total row count across all active live pipelines for auto-regeneration.
    """
    if not user_id:
        from config.settings import settings
        user_id = settings.DEFAULT_USER_ID

    try:
        total_rows = 0
        async with AsyncSessionLocal() as db:
            result = await db.execute(select(DataConnection).where(DataConnection.user_id == user_id))
            connections = result.scalars().all()
            
            # Since live row count counting across massive tables can be slow, 
            # we do a quick count heuristic or just query it if it's indexed.
            import psycopg2
            
            for conn in connections:
                if conn.source_type.lower() in ('postgres', 'postgresql'):
                    try:
                        import urllib.parse
                        safe_creds = urllib.parse.quote_plus(conn.encrypted_credentials)
                        conn_str = f"postgresql://postgres:{safe_creds}@{conn.host}/{conn.database_name}"
                        target_table = connection_target(conn)
                        if not target_table:
                            continue
                        
                        with psycopg2.connect(conn_str) as pg_conn:
                            with pg_conn.cursor() as cur:
                                # Reltuples is instant (approximate, but good enough for deltas)
                                cur.execute("SELECT reltuples::bigint FROM pg_class WHERE relname = %s", (target_table,))
                                row = cur.fetchone()
                                if row and row[0]:
                                    total_rows += row[0]
                                else:
                                    # Fallback to exact count if reltuples fails
                                    cur.execute(f"SELECT COUNT(*) FROM {target_table}")
                                    total_rows += cur.fetchone()[0]
                    except Exception as e:
                        logger.error(f"Delta check failed for connection {conn.id}: {e}")
                        
        return {"total_rows": total_rows}
    except Exception as e:
        logger.error(f"Delta endpoint error: {e}")
        return {"error": str(e)}