from fastapi import (
    FastAPI,
    HTTPException,
    WebSocket,
    WebSocketDisconnect,
    Depends,
    Header,
    Query,
    status,
)
from datetime import datetime, timedelta, timezone
from pydantic import BaseModel
import jwt
from jwt import ExpiredSignatureError, InvalidTokenError
import psycopg2
from psycopg2.extras import RealDictCursor
import asyncio


app = FastAPI()


# ============================================================
# JWT CONFIG
# ============================================================

SECRET_KEY = "my-super-secret-key-change-this"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60


# ============================================================
# REQUEST MODELS
# ============================================================

class RegisterNodeRequest(BaseModel):
    node_name: str
    host: str
    port: int = 5444
    dbname: str = "postgres"
    user: str = "postgres"
    password: str


class LinkNodeRequest(BaseModel):
    subscription_name: str

    subscriber_host: str
    subscriber_port: int = 5444

    provider_host: str
    provider_port: int = 5444

    dbname: str = "postgres"
    user: str = "postgres"
    password: str


# ============================================================
# JWT TOKEN CREATION
# ============================================================

@app.post("/spock/nodesoo")
def get_node_interfaces1():
    print("Fetching node interfaces from the database...")

    return {
        "status": "working"
    }

@app.post("/get-token")
def get_token():

    payload = {
        "sub": "frontend",
        "iat": datetime.now(timezone.utc),
        "exp": datetime.now(timezone.utc)
        + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES),
    }

    token = jwt.encode(
        payload,
        SECRET_KEY,
        algorithm=ALGORITHM,
    )

    return {
        "access_token": token,
        "token_type": "bearer",
        "expires_in": ACCESS_TOKEN_EXPIRE_MINUTES * 60,
    }


# ============================================================
# JWT VERIFICATION
# ============================================================

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",
        )


# ============================================================
# REGISTER NODE
# JWT PROTECTED
# ============================================================

@app.post("/register-node")
def register_node(
    data: RegisterNodeRequest,
    token_data: dict = Depends(verify_token),
):

    conn = None
    cur = None

    try:

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

        cur = conn.cursor()

        # Create Spock extension
        cur.execute("""
            CREATE EXTENSION IF NOT EXISTS spock;
        """)

        # Check version
        cur.execute("""
            SELECT extname, extversion
            FROM pg_extension
            WHERE extname = 'spock';
        """)

        spock = cur.fetchone()

        dsn = (
            f"host={data.host} "
            f"port={data.port} "
            f"dbname={data.dbname} "
            f"user={data.user} "
            f"password={data.password}"
        )

        # Register Spock node
        cur.execute("""
            SELECT spock.node_create(
                node_name := %s,
                dsn := %s
            );
        """, (
            data.node_name,
            dsn,
        ))

        node_id = cur.fetchone()[0]

        conn.commit()

        cur.execute("""
            SELECT node_id, node_name
            FROM spock.node;
        """)

        rows = cur.fetchall()

        nodes = [
            {
                "node_id": row[0],
                "node_name": row[1],
            }
            for row in rows
        ]

        return {
            "status": "success",
            "message": "Node registered",
            "node_id": node_id,
            "spock_version": spock[1],
            "nodes": nodes,
            "authenticated_as": token_data["sub"],
        }

    except Exception as e:

        if conn:
            conn.rollback()

        raise HTTPException(
            status_code=500,
            detail=str(e),
        )

    finally:

        if cur:
            cur.close()

        if conn:
            conn.close()


# ============================================================
# LINK NODES
# JWT PROTECTED
# ============================================================

@app.post("/link-nodes")
def link_nodes(
    data: LinkNodeRequest,
    token_data: dict = Depends(verify_token),
):

    conn = None
    cur = None

    try:

        conn = psycopg2.connect(
            host=data.subscriber_host,
            port=data.subscriber_port,
            dbname=data.dbname,
            user=data.user,
            password=data.password,
        )

        cur = conn.cursor()

        provider_dsn = (
            f"host={data.provider_host} "
            f"port={data.provider_port} "
            f"dbname={data.dbname} "
            f"user={data.user} "
            f"password={data.password}"
        )

        cur.execute("""
            SELECT spock.sub_create(
                subscription_name := %s,
                provider_dsn := %s,
                replication_sets := ARRAY[
                    'default',
                    'default_insert_only',
                    'ddl_sql'
                ],
                synchronize_structure := false,
                synchronize_data := false
            );
        """, (
            data.subscription_name,
            provider_dsn,
        ))

        subscription_id = cur.fetchone()[0]

        conn.commit()

        return {
            "status": "success",
            "message": "Nodes linked successfully",
            "subscription_id": subscription_id,
            "subscription_name": data.subscription_name,
            "subscriber": data.subscriber_host,
            "provider": data.provider_host,
            "authenticated_as": token_data["sub"],
        }

    except Exception as e:

        if conn:
            conn.rollback()

        raise HTTPException(
            status_code=500,
            detail=str(e),
        )

    finally:

        if cur:
            cur.close()

        if conn:
            conn.close()


# ============================================================
# DATABASE CONFIG
# ============================================================

DB_CONFIG = {
    "host": "172.16.10.123",
    "port": 5444,
    "dbname": "postgres",
    "user": "postgres",
    "password": "postgres",
}


# ============================================================
# SPOCK STATUS
# ============================================================

def get_spock_status():

    conn = psycopg2.connect(**DB_CONFIG)

    cursor = conn.cursor(
        cursor_factory=RealDictCursor
    )

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

    rows = cursor.fetchall()

    cursor.close()
    conn.close()

    return rows


# ============================================================
# WEBSOCKET JWT VERIFICATION
# ============================================================

def verify_websocket_token(token: str):

    try:

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

        return payload

    except ExpiredSignatureError:
        return None

    except InvalidTokenError:
        return None


# ============================================================
# WEBSOCKET
#
# ws://localhost:8000/ws/spock-status?token=JWT_TOKEN
# ============================================================

@app.websocket("/ws/spock-status")
async def websocket_spock_status(
    websocket: WebSocket,
    token: str = Query(...),
):

    token_data = verify_websocket_token(token)

    if token_data is None:

        await websocket.close(
            code=1008,
            reason="Invalid or expired JWT",
        )

        return

    await websocket.accept()

    try:

        while True:

            try:

                subscription_status = (
                    await asyncio.to_thread(
                        get_spock_status
                    )
                )

                await websocket.send_json({
                    "database_reachable": True,
                    "subscriptions": subscription_status,
                })

            except Exception as e:

                await websocket.send_json({
                    "database_reachable": False,
                    "subscriptions": [],
                    "error": str(e),
                })

            await asyncio.sleep(2)

    except WebSocketDisconnect:

        print(
            f"WebSocket client disconnected: "
            f"{token_data.get('sub')}"
        )


# ============================================================
# GET SUBSCRIPTIONS
# JWT PROTECTED
# ============================================================

@app.get("/spock/subscriptions")
def get_spock_subscription_status(
    token_data: dict = Depends(verify_token),
):

    conn = None
    cursor = None

    try:

        conn = psycopg2.connect(**DB_CONFIG)

        cursor = conn.cursor(
            cursor_factory=RealDictCursor
        )

        cursor.execute("""
            SELECT *
            FROM spock.sub_show_status();
        """)

        rows = cursor.fetchall()

        return {
            "success": True,
            "count": len(rows),
            "subscriptions": rows,
            "authenticated_as": token_data["sub"],
        }

    except Exception as e:

        raise HTTPException(
            status_code=500,
            detail=str(e),
        )

    finally:

        if cursor:
            cursor.close()

        if conn:
            conn.close()


# @app.get("/spock/node-interfaces")
# def get_node_interfaces():
#     print("Fetching node interfaces from the database...")
    # db = SessionLocal()

    # try:
    #     interfaces = db.query(NodeInterface).all()

    #     return {
    #         "status": "success",
    #         "count": len(interfaces),
    #         "nodes": [
    #             {
    #                 "if_id": node.if_id,
    #                 "if_name": node.if_name,
    #                 "if_nodeid": node.if_nodeid,
    #                 "host": node.host,
    #                 "port": node.port,
    #                 "dbname": node.dbname,
    #                 "user": node.user
    #             }
    #             for node in interfaces
    #         ]
    #     }

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

    # finally:
    #     db.close()