main.py

import os
from datetime import date, datetime, time
from decimal import Decimal

from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Query, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
import snowflake.connector

load_dotenv()

app = FastAPI(title="Snowflake Full Data API")


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # بعدين الأفضل تحط دومين الموقع بس
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


ALLOWED_TABLES = {
    "DIM_DATE",
    "DIM_PRODUCT",
    "DIM_SELLER",
    "DIM_TIME",
    "FACT_PRODUCT",
    "METRICFLOW_TIME_SPINE",
    "STG_ALL_SELLERS_PRODUCTS",
}


def verify_api_key(x_api_key: str | None = Header(default=None)):
    api_key = os.getenv("API_KEY")

    if not api_key:
        raise HTTPException(
            status_code=500,
            detail="API key is not configured on the server"
        )

    if x_api_key != api_key:
        raise HTTPException(
            status_code=401,
            detail="Invalid or missing API key"
        )

    return True


def clean_value(value):
    if isinstance(value, (datetime, date, time)):
        return value.isoformat()

    if isinstance(value, Decimal):
        return float(value)

    return value


def get_snowflake_connection():
    try:
        conn = snowflake.connector.connect(
            account=os.getenv("SNOWFLAKE_ACCOUNT"),
            user=os.getenv("SNOWFLAKE_USER"),
            password=os.getenv("SNOWFLAKE_PASSWORD"),
            role=os.getenv("SNOWFLAKE_ROLE"),
            warehouse=os.getenv("SNOWFLAKE_WAREHOUSE"),
            database=os.getenv("SNOWFLAKE_DATABASE"),
            schema=os.getenv("SNOWFLAKE_SCHEMA"),
        )
        return conn

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Snowflake connection error: {str(e)}"
        )


def run_select_query(query: str):
    conn = None
    cursor = None

    try:
        conn = get_snowflake_connection()
        cursor = conn.cursor()

        cursor.execute(query)

        columns = [col[0] for col in cursor.description]
        rows = cursor.fetchall()

        data = []

        for row in rows:
            item = {}
            for col_name, value in zip(columns, row):
                item[col_name] = clean_value(value)

            data.append(item)

        return data

    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


@app.get("/")
def home():
    return {
        "message": "Snowflake Full Data API is running",
        "note": "Protected routes need x-api-key header",
        "available_routes": [
            "/test-snowflake",
            "/tables",
            "/columns/{table_name}",
            "/data/{table_name}?limit=1000&offset=0",
            "/all-data?limit_per_table=1000&offset=0",
            "/full-database"
        ]
    }


@app.get("/test-snowflake", dependencies=[Depends(verify_api_key)])
def test_snowflake():
    conn = None
    cursor = None

    try:
        conn = get_snowflake_connection()
        cursor = conn.cursor()

        cursor.execute("SELECT CURRENT_VERSION()")
        result = cursor.fetchone()

        return {
            "status": "connected",
            "snowflake_version": result[0]
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


@app.get("/tables", dependencies=[Depends(verify_api_key)])
def get_tables():
    conn = None
    cursor = None

    try:
        conn = get_snowflake_connection()
        cursor = conn.cursor()

        cursor.execute("SHOW TABLES")
        rows = cursor.fetchall()

        tables = []

        for row in rows:
            tables.append({
                "table_name": clean_value(row[1]),
                "database": clean_value(row[2]),
                "schema": clean_value(row[3]),
                "table_type": clean_value(row[4])
            })

        return {
            "count": len(tables),
            "tables": tables
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


@app.get("/columns/{table_name}", dependencies=[Depends(verify_api_key)])
def get_columns(table_name: str):
    table_name = table_name.upper()

    if table_name not in ALLOWED_TABLES:
        raise HTTPException(status_code=400, detail="Table not allowed")

    conn = None
    cursor = None

    try:
        conn = get_snowflake_connection()
        cursor = conn.cursor()

        cursor.execute(f"DESCRIBE TABLE {table_name}")
        rows = cursor.fetchall()

        columns = []

        for row in rows:
            columns.append({
                "name": clean_value(row[0]),
                "type": clean_value(row[1]),
                "nullable": clean_value(row[3])
            })

        return {
            "table": table_name,
            "columns_count": len(columns),
            "columns": columns
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()


@app.get("/data/{table_name}", dependencies=[Depends(verify_api_key)])
def get_table_data(
    table_name: str,
    limit: int = Query(default=1000, ge=1, le=100000),
    offset: int = Query(default=0, ge=0)
):
    table_name = table_name.upper()

    if table_name not in ALLOWED_TABLES:
        raise HTTPException(status_code=400, detail="Table not allowed")

    try:
        query = f"""
            SELECT *
            FROM {table_name}
            LIMIT {limit}
            OFFSET {offset}
        """

        data = run_select_query(query)

        return {
            "table": table_name,
            "limit": limit,
            "offset": offset,
            "count": len(data),
            "data": data
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/all-data", dependencies=[Depends(verify_api_key)])
def get_all_tables_data(
    limit_per_table: int = Query(default=1000, ge=1, le=100000),
    offset: int = Query(default=0, ge=0)
):
    result = {}

    try:
        for table_name in ALLOWED_TABLES:
            query = f"""
                SELECT *
                FROM {table_name}
                LIMIT {limit_per_table}
                OFFSET {offset}
            """

            table_data = run_select_query(query)

            result[table_name] = {
                "limit": limit_per_table,
                "offset": offset,
                "count": len(table_data),
                "data": table_data
            }

        return {
            "tables_count": len(result),
            "limit_per_table": limit_per_table,
            "offset": offset,
            "tables": result
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/full-database", dependencies=[Depends(verify_api_key)])
def get_full_database():
    """
    يرجّع كل الجداول وكل الصفوف مرة واحدة.
    استخدمه فقط لو حجم الداتا مش ضخم.
    """
    result = {}

    try:
        for table_name in ALLOWED_TABLES:
            query = f"""
                SELECT *
                FROM {table_name}
            """

            table_data = run_select_query(query)

            result[table_name] = {
                "count": len(table_data),
                "data": table_data
            }

        return {
            "message": "Full database returned successfully",
            "tables_count": len(result),
            "tables": result
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    


@app.get("/count/{table_name}", dependencies=[Depends(verify_api_key)])
def get_table_count(table_name: str):
    table_name = table_name.upper()

    if table_name not in ALLOWED_TABLES:
        raise HTTPException(status_code=400, detail="Table not allowed")

    conn = None
    cursor = None

    try:
        conn = get_snowflake_connection()
        cursor = conn.cursor()

        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
        result = cursor.fetchone()

        return {
            "table": table_name,
            "total_rows": result[0]
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()