from quart import Quart, Response, request, send_from_directory, websocket
from quart_auth import (
    QuartAuth, current_user,
    AuthUser, login_user, logout_user, login_required
)
import time
import json
import asyncio
import traceback
from stream import Stream
from tool import ToolSpec, ToolParameter
from agent import AgentRun

# Vue frontend
# relative to the module file???? wtf
frontend_path = '../../frontend/dist'

def encode_json(value):
    def encode_unknown(value):
        if isinstance(value, ToolSpec):
            return value.to_openai_dict()
        elif isinstance(value, Stream):
            return {'streamId': value.id}
        else:
            return repr(value)
    return json.dumps(value, default=encode_unknown)

async def start_server(app, shutdown_event):
    server = Quart(__name__, static_folder=frontend_path)
    server.secret_key = 'horde-secret'
    auth_manager = QuartAuth(server, mode="bearer")

    stream_listeners = {}

    def cleanup(self):
        stream_listeners = []

    @server.after_request
    def after_request(response):
        response.headers.add('Access-Control-Allow-Origin', '*')
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
        response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
        return response

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

    # These paths must match the routes in the frontend SPA
    # file: frontend/src/routes/index.js
    @server.route('/user')
    @server.route('/models')
    @server.route('/remotes')
    @server.route('/runs')
    @server.route('/chat')
    @server.route('/agent')
    @server.route('/chat/<path:subpath>')
    @server.route('/agent/<path:subpath>')
    async def serve_spa_routes(subpath=None):
        # Ignore subpath and render index, Vue router will take care of everything
        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/login', methods=['POST'])
    async def login():
        data = await request.get_json()

        # token has precedence if it's present
        if "token" in data:
            user_id = auth_manager.load_token(data['token'])
        elif "username" in data and "password" in data:
            user_id = await app.verify_user_credentials(data['username'], data['password'])
        else:
            return "invalid login arguments", 400

        if user_id == None:
            return "bad login credentials", 401

        token = auth_manager.dump_token(user_id)
        login_user(AuthUser(auth_id=user_id))
        return {"token": token}

    @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 = await app.create_new_user(username, email, password)
        if user_id is None:
            return "existing or incorrect credentials", 400
        
        await app.generate_default_user_data(user_id)

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

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

        # Authenticate before dedicating any resources
        auth_req = json.loads(await websocket.receive())
        loop = asyncio.get_running_loop()

        if 'token' in auth_req:
            try:
                user_id = auth_manager.load_token(auth_req['token'])
                if not await app.verify_user_id(user_id):
                    print("Invalid Token");
                    await websocket.send(json.dumps({'error': 'invalid credentials'}))
                    return
            except Exception as err:
                print("Token authentication failed: ", err)
                await websocket.send(json.dumps({'error': 'token authentication failed'}))
                return
        else:
            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 = await app.verify_user_credentials(auth_req['username'], auth_req['password'])
            if user_id is None:
                print("Invalid Credentials");
                await websocket.send(json.dumps({'error': 'invalid credentials'}))
                return

        # TODO: print username instead, same in the finally block
        print("User logged in:", user_id)
        await websocket.send(json.dumps({'message': 'authentication successful'}))

        send_queue = asyncio.Queue()
        async def sender_impl():
            nonlocal send_queue
            while True:
                msg = await send_queue.get()
                await websocket.send(encode_json(msg))
        async def send(msg):
            await send_queue.put(msg)

        asyncio.create_task(sender_impl())
        ws_task = asyncio.current_task()
        app.ws_tasks.append(ws_task)

        session_connections = {}   # remotename -> RemoteConnection served by this socket
        stop_listening = None
        try:
            # All sessions are identical after authentication: every session
            # is a user connection, and any session can additionally connect
            # remotes ('remote-connect' action) and serve calls on them.
            app.add_user_connection(user_id, send_queue)
            
            # Listen for events and broadcast them
            async def on_event(name, data):
                await send({"action": name, "data": data})
            app.on_event(on_event)

            while True:
                req = json.loads(await websocket.receive())
                if not 'action' in req:
                    await send({'error': 'no action'})
                    continue
                action = req['action']
                body = req.get('body', {})
                nonce = req.get('nonce', None)

                async def agent_id_from_body (allow_none=False):
                    if 'agentId' in body:
                        return body['agentId']
                    elif 'agentname' in body:
                        return await app.agent_id_from_name(body['agentname'])
                    elif allow_none:
                        return None
                    else:
                        raise Exception("field 'agentId' or 'agentname' required")

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

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

                elif action == "update-agent":
                    agent_id = await agent_id_from_body()
                    await app.update_agent(
                        agent_id=agent_id,
                        name=body.get('name'),
                        agentname=body.get('agentname'),
                        description=body.get('description'),
                        picture=body.get('picture'),
                        data=body.get('data', {})
                    )
                    await send({'message': "agent updated", 'agentId': agent_id, 'nonce': nonce})

                elif action == "get-models":
                    await send({
                        'models': [{"id": model.id, "name": model.name, "value": model.value, "kind": model.kind} for model in await app.get_models()],
                        'nonce': nonce
                    })

                elif action == "save-model":
                    model = await app.save_model(
                        model_id=body.get("modelId"),
                        name=body.get("name"),
                        kind=body.get("kind"),
                        value=body.get("value"),
                    )
                    if model:
                        await send({'message': "model saved", 'modelId': model.id, 'nonce': nonce})
                    elif modelId in body:
                        await send({'message': "model deleted", 'modelId': body["modelId"], 'nonce': nonce})

                elif action == "get-agent-config-fields":
                    await send({'fields': await app.get_agent_config_fields(await agent_id_from_body()), 'nonce': nonce})

                elif action == "get-agent-kinds":
                    await send({'message': "agents retrieved", 'kinds': app.get_agent_kinds(), 'nonce': nonce})

                elif action == "get-agent-kind-info":
                    await send({'message': "agent kind retrieved", 'info': app.get_agent_kind_info(body['kind']), 'nonce': nonce})

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

                elif action == "get-agent-resources":
                    values = await app.get_agent_resources(await agent_id_from_body(), body.get('kind'), body.get('key'))
                    await send({'resources': [{'kind': r.kind, 'key': r.key, 'value': r.value} for r in values], 'nonce': nonce})

                elif action == "save-agent-resource":
                    await app.save_agent_resource(await agent_id_from_body(), body['kind'], body['key'], body['value'])
                    await send({'message': "resource saved", 'nonce': nonce})

                elif action == "get-agent-runs":
                    await send({
                        'runs': await app.get_agent_runs(agent_id=await agent_id_from_body(allow_none=True), user_id=user_id),
                        'nonce': nonce
                    })

                elif action == "cancel-agent-run":
                    await app.cancel_agent_run(body['runId'])

                elif action == "remote-list":
                    await send({'remotes': app.listRemotes(user_id), 'nonce': nonce})

                elif action == "remote-connect":
                    # Connect a remote served by THIS session's channel.
                    # Expected failures (duplicate name) are replied, not
                    # raised, so the session stays alive.
                    try:
                        connection = app.connectRemote(user_id, body['remotename'], body.get('functions', []), send)
                        session_connections[body['remotename']] = connection
                        await send({'message': "remote connected", 'remotename': body['remotename'], 'nonce': nonce})
                    except Exception as e:
                        await send({'error': str(e), 'nonce': nonce})

                elif action == "remote-disconnect":
                    # Manual disconnect, allowed from any session of the user
                    connection = app.disconnectRemote(user_id, body['remotename'])
                    session_connections.pop(body['remotename'], None)
                    if connection is None:
                        await send({'error': "no connected remote named '" + str(body['remotename']) + "'", 'nonce': nonce})
                    else:
                        await send({'message': "remote disconnected", 'remotename': body['remotename'], 'nonce': nonce})

                elif action == "remote-call":
                    # Invoke a function of one of the user's connected remotes.
                    # Expected failures (unknown remote/function, timeout,
                    # remote-side error) are replied, not raised.
                    try:
                        result = await app.callRemote(user_id, body['remotename'], body['function'], body.get('args', {}))
                        await send({'result': result, 'nonce': nonce})
                    except Exception as e:
                        await send({'error': str(e), 'nonce': nonce})

                elif action == "remote-call-result":
                    # A session serving remotes replies to a downstream call.
                    # Fields live at the message top level, not in 'body'.
                    # requestIds are '<remotename>#<n>' and resolve/reject
                    # are no-ops for unknown ids, so fanning out over this
                    # session's connections routes the reply correctly.
                    if 'error' in req:
                        for conn in session_connections.values():
                            conn.reject(req.get('requestId'), req['error'])
                    else:
                        for conn in session_connections.values():
                            conn.resolve(req.get('requestId'), req.get('result'))

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

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

                elif action == "delete-chat":
                    chat_id = await app.delete_chat(
                        chat_id=body['chatId']
                    )
                    await send({'message': "chat deleted", 'chatId': chat_id, 'nonce': nonce})

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

                elif action == "send-chat-message":
                    message = await app.send_chat_message(
                        chat_id=body['chatId'],
                        thread_id=body.get('threadId'),
                        content=body['content'],
                        kind=body.get('kind'),
                        subkind=body.get('subkind'),
                        format=body.get('format'),
                    )
                    await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})

                elif action == "update-chat-message":
                    message = await app.update_chat_message(
                        message_id=body['messageId'],
                        content=body['content'],
                        kind=body.get('kind'),
                        subkind=body.get('subkind'),
                        format=body.get('format'),
                    )
                    await send({'message': "message sent", 'messageId': body['messageId'], 'nonce': nonce})

                elif action == "create-chat-thread":
                    thread_id = await app.create_chat_thread(
                        chat_id=body['chatId'],
                        name=body.get('name'),
                    )
                    await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce})

                elif action == "send-agent-prompt":
                    if 'instruction' not in body and 'input' not in body:
                        await send({'error': "instruction or input required"})
                        continue

                    response = await app.send_agent_prompt(
                        await agent_id_from_body(),
                        body.get('instruction'),
                        body.get('input'),
                        body.get('options', {})
                    )
                    if isinstance(response, Stream):
                        if body.get('stream'):
                            response = {"streamId": response.id}
                        else:
                            response = await response.content_wait()
                    await send({'message': "prompt sent", 'response': response, 'nonce': nonce})

                elif action == "read-stream":
                    stream_id = body['streamId']
                    print("Read Stream" + str(stream_id))
                    stream = app.streams.get(stream_id)
                    if stream is None:
                        await send({'error': 'bad stream id: ' + str(stream_id), 'nonce': nonce})
                    else:
                        async def chunk_fn(chunk):
                            await send({'data': chunk, 'eof': False, 'nonce': nonce})
                        async def eof_fn():
                            await send({'eof': True, 'fullData': stream.content, 'nonce': nonce})
                        stream.listen(chunk_fn, eof_fn)

                else:
                    await send({'error': 'unknown action: ' + action, 'nonce': nonce})
        finally:
            print("User logged out:", user_id)
            if ws_task in app.ws_tasks:
                app.ws_tasks.remove(ws_task)
            # Unregister the event listener: the send closure would
            # otherwise keep queueing events for a dead socket forever
            if stop_listening:
                stop_listening()
            # All remotes served by this session are disconnected with it.
            # disconnectRemote is idempotent, so remotes that were already
            # disconnected manually (possibly from another session) are safe.
            for remotename in list(session_connections):
                app.disconnectRemote(user_id, remotename)
            app.remove_user_connection(user_id, send_queue)

    # https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
    @server.post('/api/openai/v1/chat/completions')
    @login_required
    async def openai_chat_completion():
        data = await request.get_json()

        print("Completion: " + json.dumps(data))

        # Extract required fields
        agent_id = data.get("model")
        openai_messages = data.get("messages", [])
        openai_tools = data.get("tools", [])
        stream_mode = data.get("stream", False)

        if not agent_id or not openai_messages:
            return {"error": "model and messages are required"}, 400

        user_id = current_user.auth_id
        if not user_id: pass

        # --- Parse messages ---
        message_errors = []
        messages = []
        for msg in openai_messages:
            role = msg['role']
            if role not in ['user', 'assistant', 'system']:
                message_errors.append("Bad role: " + role)
                continue
            content = msg.get('content', None)

            if type(content) == list:
                parts = content
                content = ""
                for part in parts:
                    if part['type'] == 'text':
                        content += part['text']

            if type(content) != str:
                message_errors.append("Bad content")
                continue

            messages.append({'role': role, 'content': content})
        
        if message_errors:
            print("[openai] ERROR: bad messages input\n" + "\n".join([f"\t- {e}" for e in message_errors]))

        # --- Parse tools (OpenAI format -> ToolSpec) ---
        tools = []
        for tool in openai_tools:
            if tool.get('type') != 'function':
                continue
            func = tool['function']
            params_schema = func.get('parameters', {})
            required_params = params_schema.get('required', [])
            parameters = []
            for param_name, param_spec in params_schema.get('properties', {}).items():
                parameters.append(ToolParameter(
                    name=param_name,
                    type=param_spec.get('type', 'string'),
                    description=param_spec.get('description', ''),
                    required=param_name in required_params,
                    enum=param_spec.get('enum'),
                ))
            tools.append(ToolSpec(
                name=func['name'],
                description=func.get('description', ''),
                parameters=parameters,
            ))

        # --- Get agent control and create AgentRun directly ---
        control = await app.get_agent_control(agent_id)

        run = AgentRun.create(
            app,
            agent_id=agent_id,
            user_id=user_id,
            input={
                "messages": messages,
                "tools": tools,
            },
        )
        print("run input: " + repr(run.input))

        # --- Streaming ---
        if stream_mode:
            queue = asyncio.Queue()
            loop = asyncio.get_event_loop()

            # Listen to the stream for both content and tool_call chunks

            def on_output(stream):
                async def on_chunk(chunk):
                    await queue.put(('content', chunk))
                async def on_eof():
                    # on_done doesn't work?
                    await queue.put(('eof', None))
                stream.listen(on_chunk, on_eof)

            def on_reasoning(stream):
                async def on_chunk(chunk):
                    await queue.put(('reasoning', chunk))
                async def on_eof(): pass
                stream.listen(on_chunk, on_eof)

            def on_tool_call(tool_call):
                print('agent tool call ' + repr(tool_call))
                loop.create_task(queue.put(('tool_call', tool_call)))

            def on_done(output):
                print('agent done ' + repr(output))
                loop.create_task(queue.put(('eof', None)))

            run.on('output', on_output)
            run.on('reasoning', on_reasoning)
            run.on('tool-call', on_tool_call)
            run.on('done', on_done)

            # Start the agent run
            await control.run(run)

            async def generate():
                try:
                    tool_call_count = 0
                    while True:
                        msg_type, data = await queue.get()

                        delta = None
                        if msg_type == 'content':
                            delta = {'content': data}
                        elif msg_type == 'reasoning':
                            delta = {'reasoning': data}

                        elif msg_type == 'tool_call':
                            tool_call_count = tool_call_count + 1
                            arguments = encode_json(data.arguments) if data.arguments else '{}'
                            tc_entry = {"index": tool_call_count, "function": {'name': data.tool_name, 'arguments': arguments}}
                            if data.id:
                                tc_entry['id'] = data.id
                            delta = {"tool_calls": [tc_entry]}

                        elif msg_type == 'eof':
                            # Determine finish_reason — was there a tool call?
                            finish_reason = "tool_calls" if any(
                                e['name'] == 'tool_call' for e in run.events
                            ) else "stop"

                            chunk_body = {
                                "id": "chatcmpl-" + str(time.time()),
                                "object": "chat.completion.chunk",
                                "created": int(time.time()),
                                "model": agent_id,
                                "choices": [{"delta": {}, "index": 0, "finish_reason": finish_reason}]
                            }
                            yield f"data: {json.dumps(chunk_body)}\n\n"
                            break
                        
                        if delta:
                            chunk_body = {
                                "id": "chatcmpl-" + str(time.time()),
                                "object": "chat.completion.chunk",
                                "created": int(time.time()),
                                "model": agent_id,
                                "choices": [{"delta": delta, "index": 0, "finish_reason": None}]
                            }
                            yield f"data: {json.dumps(chunk_body)}\n\n"
                        else:
                            raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")

                except Exception as e:
                    print("[openai] ERROR: " + str(e))
                    print(traceback.format_exc())
                    yield f"data: {json.dumps({'error': str(e)})}\n\n"

            return Response(generate(), content_type="text/event-stream")

        # --- Non-streaming ---
        else:
            await control.run(run)
            await run.get_output_stream().content_wait()

            # Collect any tool calls that were emitted
            tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call']

            if tool_events:
                # Aggregate tool calls: combine partial argument chunks by id
                tc_map = {}
                for evt in tool_events:
                    tid = evt.get('id', '')
                    if tid not in tc_map:
                        tc_map[tid] = {
                            "id": tid,
                            "type": "function",
                            "function": {
                                "name": evt.get('function', {}).get('name', ''),
                                "arguments": evt.get('function', {}).get('arguments', ''),
                            }
                        }
                    else:
                        tc_map[tid]['function']['arguments'] += \
                            evt.get('function', {}).get('arguments', '')

                return {
                    "id": "chatcmpl-" + str(time.time()),
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": agent_id,
                    "choices": [{
                        "message": {
                            "role": "assistant",
                            "content": None,
                            "tool_calls": list(tc_map.values()),
                        },
                        "finish_reason": "tool_calls",
                    }]
                }
            else:
                # No tool calls: return text content
                content = run.stream.content
                return {
                    "id": "chatcmpl-" + str(time.time()),
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": agent_id,
                    "choices": [{
                        "message": {
                            "role": "assistant",
                            "content": content,
                        },
                        "finish_reason": "stop",
                    }]
                }

    @server.get("/api/openai/v1/models")
    @login_required
    async def openai_models():
        """
        Return a list of available models.
        Customize this list based on your available models.
        """

        user_id = current_user.auth_id
        if not user_id: pass

        return {
            "object": "list",
            "data": [{
                "id": agent.id,
                "object": "model",
                "created": int(time.time()),
                "owned_by": user_id,
            } for agent in await app.get_agents(user_id)]
        }


    await server.run_task(debug=True, shutdown_trigger=shutdown_event.wait)