import os
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI,Depends, Header, HTTPException, status, WebSocket, WebSocketDisconnect
import psycopg2
from jwt import ExpiredSignatureError, InvalidTokenError
import jwt
from pydantic import BaseModel
from psycopg2.extras import RealDictCursor
import asyncio

load_dotenv()
app = FastAPI()

SECRET_KEY = os.getenv("SECRET_KEY")
ALGORITHM = os.getenv("ALGORITHM", "HS256")


class NodeStatusRequest(BaseModel):
    node_name: str
    host: str
    port: int
    dbname: str
    user: str
    password: str



class DatabaseRequest(BaseModel):
    host: str
    port: int
    dbname: str
    user: str
    password: str


def verify_token(
    authorization: str = Header(None),
):

    if not authorization:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Authorization header missing",
            headers={
                "WWW-Authenticate": "Bearer"
            },
        )

    parts = authorization.split()

    if len(parts) != 2:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid Authorization header",
        )

    scheme = parts[0]
    token = parts[1]

    if scheme.lower() != "bearer":
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Authorization scheme must be Bearer",
        )

    try:

        payload = jwt.decode(
            token,
            SECRET_KEY,
            algorithms=[ALGORITHM],
        )

        return payload

    except ExpiredSignatureError:

        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Token expired",
        )

    except InvalidTokenError:

        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid token",
        )
 





@app.get("/")
def root():
    return {"message": "PostgreSQL Agent is running"}



#{
#     "host": "172.16.10.178",
#     "node_name :"node3",
#     "port": 5444,
#     "dbname": "postgres",
#     "user": "postgres",
#     "password": "postgres",
#     "connect_timeout": 3,
# }


@app.get("/node/status")
def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)):

    try:
        db_config = {
            "host": data.host,
            "port": data.port,
            "dbname": data.dbname,
            "user": data.user,
            "password": data.password,
            "connect_timeout": 3,
        }
        conn = psycopg2.connect(**db_config)

        conn.close()

        return {
            "node": data.node_name,
            "online": True,
            "status": "online"
        }

    except Exception as e:

        return {
            "node": data.node_name,
            "online": False,
            "status": "offline",
            "error": str(e)
        }



latest_replication_health = None
monitor_task = None
def get_replication_health(data: NodeStatusRequest):

    db_config = {
        "host": data.host,
        "port": data.port,
        "dbname": data.dbname,
        "user": data.user,
        "password": data.password,
        "connect_timeout": 2,
        "options": "-c statement_timeout=2000"
    }

    conn = None

    try:
        conn = psycopg2.connect(**db_config)

        cur = conn.cursor(cursor_factory=RealDictCursor)

        cur.execute("""
            SELECT
                subscription_name,
                status,
                provider_node
            FROM spock.sub_show_status();
        """)

        subscriptions = cur.fetchall()

        healthy = all(
            sub["status"] == "replicating"
            for sub in subscriptions
        )

        return {
            "node": data.node_name,
            "connection_healthy": healthy,
            "subscriptions": subscriptions
        }
    except psycopg2.errors.QueryCanceled:

        print("Query cancelled because it exceeded 2 seconds")

        return {
            "status": "timeout",
            "message": "PostgreSQL query exceeded timelimit and was cancelled"
        }

    except Exception as e:

        return {
            "node": data.node_name,
            "connection_healthy": False,
            "error": str(e)
        }

    finally:
        if conn:
            conn.close()



# {
#     "node_name": "node3",
#     "host": "172.16.10.178",
#     "port": 5444,
#     "dbname": "postgres",
#     "user": "postgres",
#     "password": "postgres"
# }

@app.websocket("/ws/node/replication-health")
async def replication_monitor(websocket: WebSocket):
    await websocket.accept()
    try:
        raw_data = await websocket.receive_json()
        data = NodeStatusRequest(**raw_data)
        global latest_replication_health

        while True:

            try:

                result = await asyncio.wait_for(
                    asyncio.to_thread(
                    get_replication_health,
                    data
                    ),
                    timeout=6
                    )

                latest_replication_health = result

                print("Replication health:")
                print(result)

            except asyncio.TimeoutError:

                latest_replication_health = {
                    "node": data.node_name,
                    "connection_healthy": False,
                    "status": "timeout",
                    "message": "PostgreSQL operation exceeded maximum time"
                }

            print(latest_replication_health)
            
            await websocket.send_json(
                latest_replication_health
            )

            # Wait 2 seconds before checking again
            await asyncio.sleep(2)
    except WebSocketDisconnect:

        print("WebSocket client disconnected")



def get_database_status(
    data: DatabaseRequest
):

    conn = None

    try:
        conn = psycopg2.connect(
            host=data.host,
            port=data.port,
            dbname=data.dbname,
            user=data.user,
            password=data.password,
            connect_timeout=3
        )

        return {
            "database_name": data.dbname,
            "status": "online"
        }

    except psycopg2.Error:

        return {
            "database_name": data.dbname,
            "status": "offline"
        }

    finally:
        if conn:
            conn.close()


def get_database_uptime(
    data: DatabaseRequest
):

    conn = None

    try:
        conn = psycopg2.connect(
            host=data.host,
            port=data.port,
            dbname=data.dbname,
            user=data.user,
            password=data.password,
            connect_timeout=3
        )

        cursor = conn.cursor()

        cursor.execute("""
            SELECT EXTRACT(
                EPOCH FROM (
                    CURRENT_TIMESTAMP - pg_postmaster_start_time()
                )
            )::BIGINT;
        """)

        uptime_seconds = cursor.fetchone()[0]

        cursor.close()

        return {
            "database_name": data.dbname,
            "uptime_seconds": uptime_seconds
        }

    except psycopg2.Error as e:

        return {
            "database_name": data.dbname,
            "uptime_seconds": None,
            "error": str(e)
        }

    finally:
        if conn:
            conn.close()


@app.websocket("/ws/database/status")
async def database_status_monitor(websocket: WebSocket):

    await websocket.accept()

    try:

        raw_data = await websocket.receive_json()

        data = DatabaseRequest(**raw_data)

        while True:

            try:

                status_result, uptime_result = await asyncio.wait_for(
                    asyncio.gather(
                        asyncio.to_thread(
                            get_database_status,
                            data
                        ),
                        asyncio.to_thread(
                            get_database_uptime,
                            data
                        )
                    ),
                    timeout=5
                )

                result = {
                    "database_name": data.dbname,
                    "status": status_result["status"],
                    "uptime_seconds": uptime_result["uptime_seconds"]
                }

            except asyncio.TimeoutError:

                result = {
                    "database_name": data.dbname,
                    "status": "offline",
                    "uptime_seconds": None
                }

            await websocket.send_json(result)

            await asyncio.sleep(2)

    except WebSocketDisconnect:

        print("Database status WebSocket disconnected")


#GET ACTIVE CONNECTIONS

def get_active_connections(
    data: DatabaseRequest
):

    conn = None
    cursor = None

    try:
        conn = psycopg2.connect(
            host=data.host,
            port=data.port,
            dbname=data.dbname,
            user=data.user,
            password=data.password,
            connect_timeout=3
        )

        cursor = conn.cursor()

        cursor.execute("""
            SELECT COUNT(*)
            FROM pg_stat_activity
            WHERE state = 'active'
              AND pid <> pg_backend_pid();
        """)

        active_connections = cursor.fetchone()[0]

        return {
            "database_name": data.dbname,
            "active_connections": active_connections
        }

    except psycopg2.Error as e:

        raise HTTPException(
            status_code=500,
            detail=f"Failed to fetch active connections: {str(e)}"
        )

    finally:
        if cursor:
            cursor.close()

        if conn:
            conn.close()


# HOW MANY CONNECTIONS ARE BEING USED
def get_connection_usage(data: DatabaseRequest):

    conn = None
    cursor = None

    try:

        conn = psycopg2.connect(
            host=data.host,
            port=data.port,
            dbname=data.dbname,
            user=data.user,
            password=data.password,
            connect_timeout=3
        )

        cursor = conn.cursor()

        cursor.execute("""
            SELECT
                COUNT(*) AS current_connections,
                current_setting('max_connections')::int AS max_connections
            FROM pg_stat_activity
            GROUP BY current_setting('max_connections');
        """)

        row = cursor.fetchone()

        current_connections = row[0]
        max_connections = row[1]

        connection_usage_percent = round(
            (current_connections / max_connections) * 100,
            2
        )

        return {
            "current_connections": current_connections,
            "max_connections": max_connections,
            "connection_usage_percent": connection_usage_percent
        }

    except psycopg2.Error as e:

        return {
            "current_connections": None,
            "max_connections": None,
            "connection_usage_percent": None,
            "error": str(e)
        }

    finally:

        if cursor:
            cursor.close()

        if conn:
            conn.close()

# WEB SOCKETS FOR ACTIVE CONNECTIONS MONITORING
@app.websocket("/ws/database/active-connections")
async def database_active_connections_monitor(websocket: WebSocket):

    await websocket.accept()

    try:

        raw_data = await websocket.receive_json()

        data = DatabaseRequest(**raw_data)

        while True:

            try:

                (
                    active_connections_result,
                    connection_usage_result
                ) = await asyncio.wait_for(

                    asyncio.gather(

                        asyncio.to_thread(
                            get_active_connections,
                            data
                        ),

                        asyncio.to_thread(
                            get_connection_usage,
                            data
                        )

                    ),

                    timeout=5
                )

                result = {
                    "database_name": data.dbname,

                    "active_connections":
                        active_connections_result[
                            "active_connections"
                        ],

                    "current_connections":
                        connection_usage_result[
                            "current_connections"
                        ],

                    "max_connections":
                        connection_usage_result[
                            "max_connections"
                        ],

                    "connection_usage_percent":
                        connection_usage_result[
                            "connection_usage_percent"
                        ]
                }

            except asyncio.TimeoutError:

                result = {
                    "database_name": data.dbname,
                    "active_connections": None,
                    "current_connections": None,
                    "max_connections": None,
                    "connection_usage_percent": None
                }

            await websocket.send_json(result)

            await asyncio.sleep(2)

    except WebSocketDisconnect:

        print(
            "Database active connections WebSocket disconnected"
        )










if __name__ == "__main__":
    host = os.getenv("HOST", "127.0.0.1")
    port = int(os.getenv("PORT", 9000))

    uvicorn.run(
        "main:app",
        host=host,
        port=port,
        reload=False
    )