# -*- coding: utf-8 -*-
'''
@File    : server.py
@Time    : 2025/12/29 10:58:05
@Author  : wty-yy, Gemini3 Pro
@Version : 1.0
@Blog    : https://wty-yy.github.io/
@Desc    : Asynchronous stress pipeline evaluation server
'''
import os
os.environ['MUJOCO_GL'] = 'glfw'  # avoid mujoco.Renderer EGL context error
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

import multiprocessing
import uvicorn
import queue
import time
import uuid
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Dict, Optional

from dataclasses import dataclass
from robogauge.utils.helpers import parse_args, class_to_dict
from robogauge.tasks.pipeline.stress_pipeline import StressPipeline
from pprint import pprint

default_args_list = [
    '--stress-benchmark',
    # '--stress-terrain-names', 'flat', 'wave', 'slope', 'stairs_up', 'stairs_down', 'obstacle',
    '--stress-terrain-names', 'flat', 'wave',
    '--num-processes', '50',
    '--seeds', '0', '1', '2',
    '--search-seeds', '0', '1', '2', '3', '4',
    '--frictions', '0.5', '0.75', '1.0', '1.25', '1.5', '1.75', '2.0', '2.25', '2.5',
    '--compress-logs',
    '--headless',
]

@dataclass
class EvalTaskData:
    model_path: str
    step: int
    task_name: str
    experiment_name: str

class EvalRequest(BaseModel):
    model_path: str
    step: int
    task_name: str
    experiment_name: str

class ResponseStatus:
    PENDING = "pending"
    PROCESSING = "processing"
    FINISHED = "finished"
    ERROR = "error"
    NOT_FOUND = "not_found"

def run_api_server(input_queue, result_dict, port=9973):
    """
    Running in a separate subprocess.
    I/O Process: submit requests -> put into queue -> return ID.
    """
    app = FastAPI()

    @app.post("/submit_eval")
    def submit_eval(req: EvalRequest):
        task_id = str(uuid.uuid4())
        task_data = EvalTaskData(
            model_path=req.model_path,
            step=req.step,
            task_name=req.task_name,
            experiment_name=req.experiment_name
        )
        input_queue.put((task_id, task_data))

        result_dict[task_id] = {"status": ResponseStatus.PENDING}
        return {"task_id": task_id, "message": "Queued"}

    @app.get("/get_result/{task_id}")
    def get_result(task_id: str):
        if task_id not in result_dict:
            return {"status": ResponseStatus.NOT_FOUND}
        return result_dict[task_id]

    print(f"📡 API Server listening on port {port}...")
    uvicorn.run(app, host="127.0.0.1", port=port, log_level="error")

def main():
    print("🤖 RoboGauge Evaluation Server Starting...")
    ctx = multiprocessing.get_context('spawn')
    manager = ctx.Manager()
    task_queue = manager.Queue()
    results_store = manager.dict()

    api_p = ctx.Process(
        target=run_api_server, 
        args=(task_queue, results_store),
        daemon=True
    )
    api_p.start()

    print("🚀 Main Process started. Waiting for tasks...")
    print("   (StressPipeline will run directly in this Main Process)")

    try:
        while True:
            try:
                task_data: EvalTaskData
                task_id, task_data = task_queue.get(timeout=1.0)
                
                print(f"\n🔄 [Main] Processing Task {task_id} (Step {task_data.step})...")
                results_store[task_id] = {"status": ResponseStatus.PROCESSING}
                
                args_list = default_args_list.copy()
                args_list += ['--model-path', task_data.model_path, '--task-name', task_data.task_name, '--experiment-name', task_data.experiment_name]
                args = parse_args(args_list)

                print(f"📋 Running with args:")
                pprint(class_to_dict(args))
                
                pipeline = StressPipeline(args)
                stress_results = pipeline.run() 

                results_store[task_id] = {
                    "status": ResponseStatus.FINISHED,
                    "step": task_data.step,
                    "results": stress_results
                }
                print(f"✅ [Main] Task {task_id} Finished.")

            except queue.Empty:
                continue
            except Exception as e:
                print(f"❌ [Main] Error: {e}")
                import traceback
                traceback.print_exc()
                if 'task_id' in locals():
                    results_store[task_id] = {"status": ResponseStatus.ERROR, "error": str(e), "error_msg": traceback.format_exc()}

    except KeyboardInterrupt:
        print("\n🛑 Shutting down...")
        api_p.terminate()
        api_p.join()

if __name__ == "__main__":
    main()