script.py

import typer
from models.model import User
from database import db
import pathlib # For finding the config file, since we store it in the home dir in a hidden file
import json
from datetime import datetime
from peewee import IntegrityError, DoesNotExist
from rich.console import Console
from rich.table import Table
from rich import print as rprint

app = typer.Typer()

# NOTE: Goes through the database file when called
def init_db():
    db.connect(reuse_if_open=True)
    db.create_tables([User]) # create the models if not there

# TODO: setup the main directory
BASE_DIR = pathlib.Path.home() / ".mycli" # Setup the base folder to store the config file
CONFIG_FILE = BASE_DIR / "config.json" # This is the config file, we can use other class methods now such as unlink etc.
DB_FILE = "main.db"

# TODO: create a state object, which stores the username of the current user if the person is loggedin
class State:
    username: str | None = None

    def check_return(self) -> User | None: # First we check if they are logged in, then we check if the database file and config file exist together, then we check for an edge case where the user was removed from the table but the config file was not removed, will be handled 99% of the times by the delete-tables command
        if not self.username: # Put the logged in check here, it checks if the user is even logged in because if they are logged out, self.username will always be None
            rprint("[italic yellow]You need to be logged in to run this command")
            raise typer.Exit()
        elif not pathlib.Path(DB_FILE).exists(): # Check for the db file now that we know the session is present, if this comes true, that means the database file does not exist but the session does, this is invalid because we cant pull the info from the database without the database
            CONFIG_FILE.unlink(missing_ok=True) # again, why missing_ok? We only got to this point because the config file does exist, if the config file didn't exist, state.username would be auto set to None and the program would exit at the previous line
            rprint("[italic red]Critical error, no database file found even though a valid session exists, removing session.")
            raise typer.Exit()
        # If the code gets here, the first two checks are done, the database exists and the person is logged in 
        # run a file check at the end to see if the State instance username is actually present in the database, i dont know why we need this right now because how would the particular user be assigned to the session yet not be present in the database but because I dont know, I suppose its good measure to just implement it
        init_db()
        try:
            return User.get(username=self.username) # goes to the except block if no user, else returns the user which we can reuse now
        except DoesNotExist: # Somehow the session has an user thats not present in the database
            CONFIG_FILE.unlink()
            rprint("[italic yellow]Session has an invalid user, please relogin")
            raise typer.Exit()
    
    # NOTE: Similar as check_return but now also checks for elevated privileges, and then returns the user, or throws an error if no sufficient privileges
    def require_admin(self, host=False):
        user = self.check_return()
        if host:
            if not getattr(user, "is_host", False):
                rprint("[italic violet]This command requires host privileges")
                raise typer.Exit()
        if not getattr(user, "is_admin", False):
            rprint("[italic violet]Permission denied, need atleast admin access")
            raise typer.Exit()
            
        return user
# Initialize state for storing data
state = State()

# NOTE: typer callback decorator that runs the function everytime at the start of the app
@app.callback()
def app_load():
    # First we check the validity and if a session is present, we do nothing, just set the username class instance as the username pulled from the JSON file and continue
    if CONFIG_FILE.exists(): # check validity here
        try:
            with open(CONFIG_FILE, "r") as f:
                data =  json.load(f) # Load the file, will throw an error and go to except if the file is of invalid json
            if isinstance(data, dict) and data.get("username"):
                # we get here if the file is completely valid
                # now we just take the username from the session folder and keep it on state so that we dont need to reopen the JSON file
                state.username = data["username"]
            else:
                # Either its not a dict or you cant get username (someone just removed the username), so we remove the file 
                CONFIG_FILE.unlink()
        except json.JSONDecodeError:
            CONFIG_FILE.unlink() # after opening, if the json file is invalid, then remove the json file
    else:
        pass
        # Get here if the file does not exist which means they are most probably not logged in, we dont do anything here

# OPTIMIZE: login
@app.command()
def login(name: str, number: int):
    init_db()
    if not state.username: # no user has logged in so username is now None
        try:
            user = User.get(username=name, number=number)

            # we get the user, its gonna throw an error if the user does not exist and skip to the except block, else now we know the user exists
            # just log him in now and then save to session later

            user.time_logged_in = datetime.now()
            user.save()
            # close the db connection and save changes
            rprint("[italic green]Found user, updated login time")
        except DoesNotExist:
            try:
                rprint("[italic yellow]Couldn't find user, do you want me to create a new user?")
                confirm = typer.confirm("")
                if confirm:
                    User.create(username=name, number=number, time_logged_in=datetime.now())
                else:
                    rprint("[italic red]Exiting....")
                    raise typer.Exit()
            except IntegrityError:
                typer.echo("[italic red]User already exists, please enter the correct password")
                raise typer.Exit()
            # we create the base User
        # Since both are done, now we just save the user to the session, because either way, we are logging in now.
        # first create the directory because we get here since we never logged in, we should not have the file
        BASE_DIR.mkdir(exist_ok=True)
        with open(CONFIG_FILE, "w") as f:
            json.dump({"username": name}, f)
        rprint("[italic green]Welcome {0}".format(name))
    else:
        rprint("[italic yellow]You are already logged in as {0}".format(state.username))
        raise typer.Exit()

# NOTE: logout
@app.command()
def logout(): # No need for arguments, we pull the name from the sessions folder, and if it does not exist we just exit
    user: User = state.check_return() # type: ignore
    CONFIG_FILE.unlink() # I dont need to update and reupdate the state because once they rerun the file, the state will be reset either way and it wont assign the old username because there is no username to assign
    rprint("[italic yellow]Goodbye {0}".format(user.username))
    # the old state is accessible to this line


# NOTE: testcase, basic command to check if the session logic is working
@app.command()
def whoami():
    user: User = state.check_return() # type: ignore
    # Make a table using rich
    table = Table(title="Account Details")
    table.add_column("Type")
    table.add_column("Info")
    table.add_row("Username", str(user.username))
    table.add_row("Password", str(user.number))
    table.add_row("Account created", user.account_created.strftime("%B %d, %Y at %I:%M %p")) # or just use the user.account_created.year format although that is far longer lol
    table.add_row("Time logged in", user.time_logged_in.strftime("%B %d, %Y at %I:%M %p"))
    if user.is_admin and (not user.is_host):
        table.add_row("Admin", "Yes")
    if user.is_host:
        table.add_row("Host", "Yes")
    console = Console()
    console.print(table)

# TODO: Delete command first runs and checks if the user is logged in, prompts them again if they really do want to, and then deletes the account, also wipes it from session
@app.command()
def delete():
    user: User = state.check_return() # type: ignore
    confirmation = typer.confirm("Are you sure? ") # Keep it concrete, just ignore the case
    if not confirmation:
        raise typer.Abort()
    CONFIG_FILE.unlink()
    f_u = User.get_or_none(username=user.username)
    if not f_u:
        typer.echo("Critical error: Bypassed state check return yet user was not present in the database")
        typer.Exit()
    f_u.delete_instance()
    typer.echo("Your account {0} was succesfully deleted".format(user.username))
    raise typer.Exit()       

#TODO: Delete-user: Requires admin privileges. Removes the username provided from the table and the session(actually kind of useless because only one session exists at a time and if i remove session file ill end up removing the admins session lol) but only if the user who has the isAdmin set to true
@app.command()
def delete_user():
    state.require_admin()
    # We now know the user has admin privileges
    # ask for confirmation
    conf1 = typer.confirm("Do you want to continue with this command")
    if not conf1:
        raise typer.Abort()
    while True:
        prompt1 = typer.prompt("User", type=str)
        user = User.get_or_none(username=prompt1)
        if not user:
            rprint("[italic yellow]Invalid user")
            continue
        if not state.username == prompt1:
            conf2 = typer.confirm("[italic orange]Found user, do you want to continue?")
            if not conf2:
                break
            user.delete_instance()
            typer.echo("[italic green]Deleted user")
            raise typer.Exit()
        else:
            print("You can't delete yourself!")
            continue

# TODO: initial init command that allows the user to create the first user with admin privileges, only works once.
@app.command(hidden=True)
def init(name: str, number: int):
    init_db()

    if User.select().count() > 0: # If more than zero users, refuse usage of this command
        typer.echo("Command use has been exceeded")
        raise typer.Exit()
    
    User.create(username=name, 
                number=number, 
                time_logged_in=datetime.now(),
                is_admin=True,
                is_host=True
                )
    # We create the user, then we check for the BASE_DIR, create it or not if it already exists
    # Then we open the config_file and save the session
    # New user now present with host permissions, it never reaches here if there are accounts
    # It throws an error an exits
    BASE_DIR.mkdir(exist_ok=True)
    with open(CONFIG_FILE, "w") as f:
        json.dump({"username": name}, f)
    print("Created user {0} with host privileges and logged in".format(name))

#TODO: shorten the creation of user and session by creating a class method in the state class called create and save that takes the arguments from the user, and then writes down the session saving and database code itself, to prevent redundancy and ease readability


#TODO: is_admin uses the set-admin command, which can only be run by the host, add host_return() and admin_return() commands in hierarchy for host and admin checks, however i need a better way to check for the host, this is alaso for future implementation
@app.command()
def set_admin(name:str):
    state.require_admin(host=True) # type: ignore
    # User now has host if they reach past this
    target_user = User.get_or_none(username=name)
    
    if not target_user:
        rprint("[italic red]Invalid user.")
    if target_user.is_admin:
        rprint("[italic red]Already is an admin.")
        typer.echo("")
        raise typer.Exit()
    
    rprint("[italic yellow]Are you sure?\nThis will allow the user to run administrator commands which can be harmful without proper overlook")
    confirm = typer.confirm("")
    if confirm:
        target_user.is_admin = True
        target_user.save()
        rprint("[italic green]Gave admin privileges to {0}".format(name))
    else:
        rprint("[italic red]Exiting...")

@app.command()
def user_info(name: str):
    ...


#TODO: finally implement the main features such as set-tasks, remove-tasks, edit-tasks etc. as a normal task manager app does
#BUG: The code might break if the user somehow figures out a way to remove admin privileges but keep host privileges. Then the host wont be able to use admin commands which is just wrong
#TODO: Make a way to show the date and year and other parts of the date time command in a readable human formatish
#TODO: Update the whomai command to show host: yes if true or admin: yes if true and show nothing if neither host nor admin
#OPTIMIZE: Create a modular file system where all the commands are auto registered by taking it from a folder
#OPTIMIZE: Would it not be more user friendly if the app told the user what permission it needed to run the commands, instead of prompting it admin permissions, and then host permissions
#OPTIMIZE: any way to make the code run faster?


if __name__=="__main__":
    app()