from fastapi import FastAPI, Depends, HTTPException, Response, Request, status, Cookie
from psycopg_pool import AsyncConnectionPool
import psycopg
import uvicorn
import os
from pydantic import BaseModel, Field, model_validator, field_validator
import time
from datetime import date, datetime, timedelta, timezone
from confluent_kafka import Producer
import uuid
import hashlib
import jwt

# class Item(BaseModel):
#     name: str
#     description: str
#     reporter: str
#     priority: int
#     is_stupid: bool


with open('keys/ec_private.pem', 'rb') as f:
    PRIVATE_KEY = f.read()

with open('keys/ec_public.pem', 'rb') as f:
    PUBLIC_KEY = f.read()

REGISTRATION_MINIMUM_AGE = 4745

def get_month_length(month, year = time.localtime().tm_year):
    assert not (month > 12 or month < 1)
    return {
        1: 31,
        2: 29 if year % 4 == 0 else 28,
        3: 31,
        4: 30,
        5: 31,
        6: 30,
        7: 31,
        8: 31,
        9: 30,
        10: 31,
        11: 30,
        12: 31
    }[month]

class User(BaseModel):
    email: str = Field(..., max_length = 64)
    password: str
    birthYear: int
    birthMonth: int = Field(..., ge=1, le=12)
    birthDay: int
    firstName: str = Field(..., max_length = 64)
    lastName: str = Field(..., max_length = 64)
    preInstalledApplications: dict
    
    @model_validator(mode='before')
    def strip_all(cls, values):
        return {k: v.strip() if type(v) and k != 'password' == str else v for k, v in values.items()}

    @field_validator('email')
    def validate_email(cls, value):
        def e(): raise ValueError('invalid e-mail address')
        if not '@' in value or ' ' in value:
            e()
        user, host = value.split('@')
        if host.endswith('.'):
            host = host[:-1]
        if not user or not host or '.' not in host or host.startswith('.'):
            e()
        labels = host.split('.')
        for l in labels:
            if len(l.encode()) < 1 or len(l.encode()) > 63:
                e()
        return value

    @field_validator('password')
    def validate_password_security(cls, password):
        password = password
        number_of_caps     = 0
        number_of_smalls   = 0
        number_of_numbers  = 0
        number_of_specials = 0
        for char in password:
            if char.isalnum():
                if char.isalpha():
                    if char.capitalize() == char:
                        number_of_caps += 1
                    else:
                        number_of_smalls += 1
                else:
                    number_of_numbers += 1
            else:
                number_of_specials += 1
        if not (number_of_caps > 0 and number_of_numbers > 0 and number_of_smalls > 0 and number_of_specials > 0 and len(password) > 7):
            raise ValueError('the password does not match the requirements')
        return password
    
    @model_validator(mode='before')
    def validate_pre_installed_applications(cls, values):
        ...
        return values
    
    @model_validator(mode='before')
    def validate_birth_date(cls, values):
        day = values.get('birthDay')
        if day < 1 or day > get_month_length(values.get('birthMonth'), values.get('birthYear')):
            raise ValueError('invalid birth day')
        today = date.today()
        birth_date = date(values.get('birthYear'), values.get('birthMonth'), day)
        if not (today >= birth_date):
            raise ValueError('the date cannot be in future')
        if (today - birth_date).days < REGISTRATION_MINIMUM_AGE:
            raise ValueError('age restriction, minimum age is 13 years')
        return values


class Session(BaseModel):
    password: str
    stayLoggedIn: bool = Field(False)

app = FastAPI(
    # docs_url = None,
    # redoc_url = None,
    # openapi_url = None
)
pool: AsyncConnectionPool

@app.on_event('startup')
async def on_startup():
    global pool
    pool = AsyncConnectionPool(
        os.getenv('DB_URL'),
        min_size=5,
        max_size=40
    )
    await pool.open()
    print('[INFO] Started connection pool')

    async with pool.connection() as conn:
        async with conn.cursor() as cursor:
            await cursor.execute('CREATE TABLE IF NOT EXISTS users (id UUID PRIMARY KEY, email VARCHAR(64) UNIQUE, phone VARCHAR(32), password BYTEA)')
            await cursor.execute('CREATE TABLE IF NOT EXISTS recovery (id UUID PRIMARY KEY, remail VARCHAR(64))')
            print('[INFO] Database initialized')
        await conn.commit()

@app.on_event('shutdown')
async def on_shutdown():
    await pool.close()
    print('[INFO] Closed connection pool')

async def get_conn():
    async with pool.connection() as conn:
        yield conn

@app.post('/users')
async def create_user(user: User, conn = Depends(get_conn)):
    '''Register a user'''
    async with conn.cursor() as cursor:
        user_id = uuid.uuid4()
        print(await (await cursor.execute('SELECT * FROM users')).fetchall())
        print(await (await cursor.execute('SELECT * FROM recovery')).fetchall())
        try:
            await cursor.execute('INSERT INTO users (id, email, phone, password) VALUES (%s, %s, %s, %s)', (user_id, user.email, None, hashlib.sha3_512(user.password.encode()).digest()))
            await cursor.execute('INSERT INTO recovery (id, remail) VALUES (%s, %s)', (user_id, None))
        except psycopg.errors.UniqueViolation: 
            raise HTTPException(status_code=409, detail = 'this user is already registered')
        conn.commit()
    return {
        'id': user_id,
        'email': user.email,
        'remail': None,
        'phone': None
    }

@app.post('/users/{user}/session')
async def create_session(response: Response, session: Session, user: uuid.UUID, conn = Depends(get_conn)):
    '''Create a session'''
    async with conn.cursor() as cursor:
        await cursor.execute('SELECT COUNT(*) FROM users WHERE id = %s', (user,))
        if not (await cursor.fetchone())[0]:
            raise HTTPException(status_code = 404, detail = 'user not found')
        await cursor.execute('SELECT password FROM users WHERE id = %s', (user,))
        if hashlib.sha3_512(session.password.encode()).digest() != (await cursor.fetchone())[0]:
            raise HTTPException(status_code = 401, detail = 'Unauthorized')
        payload = {
            'iss': 'http://localhost:1234',
            'sub': str(user),
            'aud': 'jcloud-auth',
            'iat': datetime.now(tz=timezone.utc),
            'exp': datetime.now(tz=timezone.utc) + timedelta(days=20),
        }
        token = jwt.encode(
            payload,
            PRIVATE_KEY,
            algorithm = 'ES256',
            # headers = {
            #     'kid': 'ec-key-2025-01'
            # }
        )
        response.set_cookie(
            key = 'auth_token',
            value = token,
            max_age = 1728000 if session.stayLoggedIn else None,
            httponly = True,
            secure = False,
            samesite = 'strict'
        )
        response.set_cookie(
            key = 'user',
            value = str(user),
            max_age = 1728000 if session.stayLoggedIn else None,
            httponly = True,
            secure = False,
            samesite = 'strict',
        )
        response.set_cookie
        return {
            'user': user,
            'authToken': token,
            'exp': datetime.now(tz=timezone.utc) + timedelta(days=20)
        }

@app.delete('/users/session', status_code = status.HTTP_204_NO_CONTENT)
async def delete_session(response: Response):
    '''Delete a session'''
    response.set_cookie(
        key = 'auth_token',
        value = '',
        max_age = -1,
        httponly = True,
        secure = False
    )
    response.set_cookie(
        key = 'user',
        value = '',
        max_age = -1,
        httponly = True,
        secure = False
    )
    
@app.get('/users/{user}/validity')
async def validate_token(user: uuid.UUID, conn = Depends(get_conn), Auth: str = Cookie(...)):
    '''Validate a token'''
    async with conn.cursor() as cursor:
        await cursor.execute('SELECT COUNT(*) FROM users WHERE id = %s', (user,))
        if not (await cursor.fetchone())[0]:
            raise HTTPException(status_code = 404, detail = 'User not found')
        try:
            jwt.decode(
                Auth,
                PUBLIC_KEY,
                algorithms = ['ES256'],
                audience = 'jcloud-auth',
                issuer = 'http://localhost:1234',
                subject = str(user)
            )
        except:
            raise HTTPException(status_code = 401, detail = 'Unauthorized')
        return b'success'
    

@app.get('/health')
async def get_health_status():
    '''Sends `OK`'''
    return b'OK'

# @app.post('/item')
# async def create_item(item: Item, conn = Depends(get_conn)):
#     async with conn.cursor() as cursor:
#         await cursor.execute('INSERT INTO items (name, description, reporter, priority, is_stupid) VALUES (%s, %s, %s, %s, %s) RETURNING *', (item.name, item.description, item.reporter, item.priority, item.is_stupid))
#         i = await cursor.fetchone()
#         return {'id': i[0], 'description' : i[1], 'reporter': i[2], 'priority': i[3], 'is_stupid': i[4]}

# @app.get('/item/{item_id}')
# async def read_item(item_id, conn = Depends(get_conn)):
#     async with conn.cursor() as cursor:
#         await cursor.execute('SELECT * FROM items WHERE id = %s', (item_id,))
#         i = await cursor.fetchone()
#         return {'id': i[0], 'name' : i[1], 'description': i[2], 'reporter': i[3], 'priority': i[4], 'is_stupid': i[5]}

# @app.get('/items')
# async def read_item(conn = Depends(get_conn)):
#     async with conn.cursor() as cursor:
#         await cursor.execute('SELECT * FROM items')
#         items = await cursor.fetchall()
#         return [{'id': i[0], 'description' : i[1], 'reporter': i[2], 'priority': i[3], 'is_stupid': i[4]} for i in items]

# @app.get('/datadir')
# async def get_datadir(conn = Depends(get_conn)):
#     async with conn.cursor() as cursor:
#         await cursor.execute('SHOW data_directory')
#         return {'res': await cursor.fetchall()}

if __name__ == '__main__':
    uvicorn.run(
        'main:app',
        host = '0.0.0.0',
        port = 1234,
        reload = True,
        reload_dirs = ['/app'],
        server_header = False
    )