from datetime import datetime
import asyncio
import json
import sys
import os

from sqlalchemy.orm import Session
from sqlalchemy.sql.expression import select
from quart import Quart, Response, request, send_from_directory, jsonify, websocket
import replicate

import database as db
from server import start_server
from agent import Agent

MESSAGE_PAGE_SIZE = 32

class ChatContext:
    def __init__(self, messages, prelude=None):
        self.messages=messages
        self.prelude=prelude

class App:
    def __init__(self, data_directory, replicate_api_token=None):
        self.data_directory = data_directory
        self.sql_engine = db.create_sql_engine(self.data_directory)
        self.user_connections = {}
        self.remote_connections = {}
        self.streams = {}

        if replicate_api_token:
            self.replicate_client = replicate.Client(api_token=replicate_api_token)
        else:
            raise Exception("Horde only works with a replicate client. Please provide an API Token.")

        # `Agent` is the kind 'basic', will be class `BasicAgent` in the future
        self.agent_classes = [Agent]

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

    def add_user_connection(self, user_id, channel):
        if user_id not in self.user_connections:
            self.user_connections[user_id] = []
        if channel not in self.user_connections[user_id]:
            self.user_connections[user_id].append(channel)

    def remove_user_connection(self, user_id, channel):
        if user_id in self.user_connections:
            if channel in self.user_connections[user_id]:
                self.user_connections[user_id].remove(channel)
                if not self.user_connections[user_id]:  # Check if the list is empty
                    del self.user_connections[user_id]

    def add_remote_connection(self, user_id, remotename, channel):
        if user_id not in self.remote_connections:
            self.remote_connections[user_id] = {}
        user_connections = self.remote_connections[user_id]

        if channel_list not in user_connections[remotename]:
            user_connections[remotename] = []
        user_connections[remotename].append(channel)

    def create_agent(self, user_id, name, description, picture, kind, data):
        with self.sql_session() as session:
            agent = db.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, members=[]):
        with self.sql_session() as session:
            chat = db.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}")

            chat_id = chat.id
        
        for agent_id in members:
            self.add_agent_to_chat(chat_id, agent_id)

        return chat_id

    # Updates the agent's basic fields, and applies field updates in its data
    def update_agent(self, agent_id, name=None, description=None, picture=None, data=None):
        with self.sql_session() as session:
            agent = session.execute(
                select(db.Agent).where(db.Agent.id == agent_id)
            ).scalars().first()

            if not agent:
                raise Exception(f"Agent not found with id {agent_id}")
            
            summary = []

            if name != None:
                agent.name = name
                summary.append('name')
            if description != None:
                agent.description = description
                summary.append('description')
            if picture != None:
                agent.picture = picture
                summary.append('picture')

            if data:
                # Assign a copy to trigger a value change
                agent.data = agent.data.copy()
                agent.data.update(data)
                summary.extend([f"data.{name}" for name in data.keys()])

            session.commit()
            print(f"Updated agent #{agent.id} fields: {summary}")

    def get_chats(self, user_id):
        with self.sql_session() as session:
            chats = session.execute(select(db.Chat).where(db.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(db.Agent).where(db.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_agent_kinds(self):
        return [kind.identifier for kind in self.agent_classes]

    def delete_chat(self, chat_id):
        with self.sql_session() as session:
            session.execute(db.Membership.__table__.delete().where(db.Membership.chat_id == chat_id))
            session.execute(db.Chat.__table__.delete().where(db.Chat.id == chat_id))
            session.commit()
            print(f"Deleted chat #{chat_id} and its memberships.")

    def get_chat_messages(self, chat_id, earliest_index=0):
        with self.sql_session() as session:
            messages = session.execute(
                select(db.Message).where(db.Message.chat_id == chat_id).order_by(db.Message.time.asc())
            ).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(db.User).where(db.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(db.User.id).where(db.User.email == email or db.User.username == username)
            ).first()

            if result is None:
                user = db.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 = db.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}")

    # If an integer is passed in sender, it is interpreted as an agent id
    def send_chat_message(self, chat_id, content, sender='user'):
        with self.sql_session() as session:
            message = db.Message()
            message.chat_id = chat_id
            message.time = int(datetime.now().timestamp())
            message.sender = sender
            message.kind = 'text'  # default kind
            message.content = content

            session.add(message)
            session.commit()

            # Notify chat owner
            if sender != 'user':
                chat_owner = session.execute(select(db.Chat.user_id).where(db.Chat.id == chat_id)).scalar()
                self.notify_message_to_user(chat_owner, message)

            # Notify participant agents
            agent_ids = session.execute(select(db.Membership.agent_id).where(db.Membership.chat_id == chat_id)).scalars().all()
            for agent_id in agent_ids:
                if agent_id != sender:
                    self.notify_message_to_agent(agent_id, message)

            return {
                'id': message.id,
                'chatId': message.chat_id,
                'time': message.time,
                'sender': message.sender,
                'kind': message.kind,
                'content': message.content,
            }

    def get_chat_context(self, chat_id):
        return ChatContext(self.get_chat_messages(chat_id))

    def get_agent_control(self, agent_id):
        with self.sql_session() as session:
            agent = session.execute(
                select(db.Agent).where(db.Agent.id == agent_id)
            ).scalars().first()

            if not agent:
                raise Exception(f"Agent not found with id {agent_id}")
        
            control = Agent.create_from_data(app, agent_id, agent.kind, agent.data)
            return control

    def get_agent_config_fields(self, agent_id):
        with self.sql_session() as session:
            agent = session.execute(
                select(db.Agent).where(db.Agent.id == agent_id)
            ).scalars().first()
            if not agent:
                raise Exception(f"Agent not found with id {agent_id}")

            agent_class = next((it for it in self.agent_classes if it.identifier == agent.kind), None)
            if not agent_class:
                raise Exception(f"Invalid agent kind: {agent.kind}")

            data = agent.data or {}
            fields = [{
                'name': field.name,
                'type': field.type,
                'value': data.get(field.name, field.default)
                } for field in agent_class.schema_fields]
            
            return fields
        return {}

    async def send_agent_prompt(self, agent_id, message):
        control = self.get_agent_control(agent_id)
        response = await control.prompt(message)
        return response

    async def call_agent_action(self, agent_id, action, args):
        control = self.get_agent_control(agent_id)
        response = await control.execute_action(action, args)
        return response

    def notify_message_to_user(self, user_id, message):
        if user_id in self.user_connections:
            user_message = {
                'action': 'chat-message',
                'id': message.id,
                'chatId': message.chat_id,
                'time': message.time,
                'sender': message.sender,
                'kind': message.kind,
                'content': message.content
            }
            for channel in self.user_connections[user_id]:
                asyncio.create_task(channel.put(user_message))

    def notify_message_to_agent(self, agent_id, message):
        control = self.get_agent_control(agent_id)
        asyncio.create_task(control.received_message(message))

    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="basic", # Should be '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,
            members=[agent_id],
        )

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 = "."

    print("api token", os.getenv('REPLICATE_API_TOKEN'))

    app = App(
        data_directory=data_directory,
        replicate_api_token=os.getenv('REPLICATE_API_TOKEN'),
    )
    start_server(app)