import argparse
import json
import logging
import time

import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.requests import Request

from lagent.schema import AgentMessage
from lagent.utils import load_class_from_string


class AgentAPIServer:

    def __init__(self,
                 config: dict,
                 host: str = '127.0.0.1',
                 port: int = 8090):
        self.app = FastAPI(docs_url='/')
        self.app.add_middleware(
            CORSMiddleware,
            allow_origins=['*'],
            allow_credentials=True,
            allow_methods=['*'],
            allow_headers=['*'],
        )
        cls_name = config.pop('type')
        python_path = config.pop('python_path', None)
        cls_name = load_class_from_string(cls_name, python_path) if isinstance(
            cls_name, str) else cls_name
        self.agent = cls_name(**config)
        self.setup_routes()
        self.run(host, port)

    def setup_routes(self):

        def heartbeat():
            return {'status': 'success', 'timestamp': time.time()}

        async def process_message(request: Request):
            try:
                body = await request.json()
                message = [
                    m if isinstance(m, str) else AgentMessage.model_validate(m)
                    for m in body.pop('message')
                ]
                result = await self.agent(*message, **body)
                return result
            except Exception as e:
                logging.error(f'Error processing message: {str(e)}')
                raise HTTPException(
                    status_code=500, detail='Internal Server Error')

        def get_memory(session_id: int = 0):
            try:
                result = self.agent.state_dict(session_id)
                return result
            except KeyError:
                raise HTTPException(
                    status_code=404, detail="Session ID not found")
            except Exception as e:
                logging.error(f'Error processing message: {str(e)}')
                raise HTTPException(
                    status_code=500, detail='Internal Server Error')

        self.app.add_api_route('/health_check', heartbeat, methods=['GET'])
        self.app.add_api_route(
            '/chat_completion', process_message, methods=['POST'])
        self.app.add_api_route(
            '/memory/{session_id}', get_memory, methods=['GET'])

    def run(self, host='127.0.0.1', port=8090):
        logging.info(f'Starting server at {host}:{port}')
        uvicorn.run(self.app, host=host, port=port)


def parse_args():
    parser = argparse.ArgumentParser(description='Async Agent API Server')
    parser.add_argument('--host', type=str, default='127.0.0.1')
    parser.add_argument('--port', type=int, default=8090)
    parser.add_argument(
        '--config',
        type=json.loads,
        required=True,
        help='JSON configuration for the agent')
    args = parser.parse_args()

    return args


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    args = parse_args()
    AgentAPIServer(args.config, host=args.host, port=args.port)