from database import create_sql_engine, User, Agent, Chat, Message, Membership
from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import select
from quart import Quart, Response, request, send_from_directory, jsonify, websocket
import asyncio
import json
import sys

class App:
    def __init__(self, data_directory):
        self.data_directory = data_directory
        self.sql_engine = create_sql_engine(self.data_directory)

    def sql_session(self):
        return Session(self.sql_engine)

    def create_agent(self, user_id, name, description, picture, kind, data):
        with self.sql_session() as session:
            agent = Agent()
            agent.user_id = user_id
            agent.name = name
            agent.description = description
            agent.picture = picture
            agent.kind = kind
            agent.data = data

            session.add(agent)
            session.commit()
            print(f"Created agent #{agent.id}: {agent.name}")
            return agent.id

    def create_chat(self, user_id, name, description, picture):
        with self.sql_session() as session:
            chat = Chat()
            chat.user_id = user_id
            chat.name = name
            chat.description = description
            chat.picture = picture

            session.add(chat)
            session.commit()
            print(f"Created chat #{chat.id}: {chat.name}")
            return chat.id

    def get_chats(self, user_id):
        with self.sql_session() as session:
            chats = session.execute(select(Chat).where(Chat.user_id == user_id)).scalars().all()
            return [{
                "id": chat.id,
                "name": chat.name,
                "description": chat.description,
                "picture": chat.picture
            } for chat in chats]

    def get_agents(self, user_id):
        with self.sql_session() as session:
            agents = session.execute(select(Agent).where(Agent.user_id == user_id)).scalars().all()
            return [{
                "id": agent.id,
                "name": agent.name,
                "description": agent.description,
                "picture": agent.picture
            } for agent in agents]

    def get_chat_messages(self, chat_id, earliest_index=0):
        with self.sql_session() as session:
            messages = session.execute(
                select(Message).where(Message.chat_id == chat_id).order_by(Message.id.desc())
            ).scalars().fetchmany(MESSAGE_PAGE_SIZE)

            return [{
                "id": message.id,
                "time": message.time,
                "sender": message.sender,
                "kind": message.kind,
                "content": message.content
            } for message in messages]

    def verify_user_credentials(self, username, password):
        with self.sql_session() as session:
            user = session.execute(
                select(User).where(User.username == username)
            ).scalars().first()

            if user and user.password == password:
                return user.id
            return None

    def create_new_user(self, username, email, password):
        with self.sql_session() as session:
            result = session.execute(
                select(User.id).where(User.email == email or User.username == username)
            ).first()

            if result is None:
                user = User()
                user.email = email
                user.username = username
                user.password = password

                session.add(user)
                session.commit()
                
                user_id = user.id
                print(f"Created user #{user_id}: {username} ({email})")
            else: 
                print(f"Signup error: '{email}' or '{username}' are already used")
        return user_id
    
    def add_agent_to_chat(self, chat_id, agent_id):
        with self.sql_session() as session:
            membership = Membership()
            membership.chat_id = chat_id
            membership.agent_id = agent_id
            
            session.add(membership)
            session.commit()
            print(f"Added agent #{agent_id} to chat #{chat_id}")

    def generate_default_user_data(self, user_id):
        agent_id = self.create_agent(
            user_id=user_id,
            name="Intro",
            description="Introduces you to the application.",
            picture=None,
            kind="intro",
            data={}
        )

        # Create a chat with that agent in it
        chat_id = self.create_chat(
            user_id=user_id,
            name="Initial Chat",
            description=None,
            picture=None
        )

        # Add the agent to the chat
        self.add_agent_to_chat(chat_id, agent_id)

def start_server(app):
    server = Quart(__name__, static_folder='../frontend/build/web')

    @server.route('/')
    async def index():
        return await send_from_directory(server.static_folder, 'index.html')

    @server.route('/<path:path>')
    async def serve_static(path):
        return await send_from_directory(server.static_folder, path)

    @server.route('/user/signup', methods=['POST'])
    async def signup():
        data = await request.get_json()

        user_id = None

        username = data['username']
        email = data['email']
        password = data['password']

        user_id = app.create_new_user(username, email, password)
        if user_id is None:
            return {"error": "existing credentials"}
        
        app.generate_default_user_data(user_id)

        # token = auth_manager.dump_token(user_id)
        return {"userId": user_id}

    @server.websocket('/ws')
    async def ws():
        print("incoming ws")

        # Authenticate before dedicating any resources
        auth_req = json.loads(await websocket.receive())
        if "username" not in auth_req or "password" not in auth_req:
            print("Invalid auth request:", auth_req)
            await websocket.send(json.dumps({'error': 'invalid auth'}))
            return

        user_id = app.verify_user_credentials(auth_req['username'], auth_req['password'])
        if user_id is None:
            await websocket.send(json.dumps({'error': 'invalid credentials'}))
            return

        send_queue = asyncio.Queue()

        async def sender_task():
            nonlocal send_queue
            while True:
                msg = await send_queue.get()
                await websocket.send(json.dumps(msg))

        async def send(msg):
            await send_queue.put(msg)

        asyncio.create_task(sender_task())

        while True:
            req = json.loads(await websocket.receive())
            action = req['action']
            body = req.get('body', {})
            nonce = req.get('nonce', None)

            if action == "ping":
                print("ping")
                await send({'message': "pong"})

            elif action == "create-agent":
                agent_id = app.create_agent(
                    user_id=user_id,
                    name=body['name'],
                    description=body['description'],
                    picture=body.get('picture'),
                    kind=body['kind'],
                    data=body['data']
                )
                await send({'message': "agent created", 'agentId': agent_id, 'nonce': nonce})

            elif action == "get-agents":
                agents_data = app.get_agents(user_id)
                await send({'message': "agents retrieved", 'agents': agents_data, 'nonce': nonce})

            elif action == "create-chat":
                chat_id = app.create_chat(
                    user_id=user_id,
                    name=body['name'],
                    description=body.get('description'),
                    picture=body.get('picture')
                )
                await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce})

            elif action == "get-chats":
                chats_data = app.get_chats(user_id)
                await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce})

            elif action == "get-chat-messages":
                messages_data = app.get_chat_messages(
                    body['chatId'],
                    earliest_index=body.get('earliestIndex', 0)
                )
                await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce})

            else:
                print("Unknown message: " + action)

    server.run(debug=True)

if __name__ == "__main__":
    if len(sys.argv) > 1:
        data_directory = sys.argv[1]
    else:
        print("No directory specified. Using current directory as default.")
        data_directory = "."

    app = App(data_directory)
    start_server(app)