#  Copyright 2026 zhaoxi826
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.


from typing import Union
from pretor.utils.ray_hook import ray_actor_hook
from fastapi import APIRouter, Depends
from pydantic import BaseModel
from pretor.utils.access import Accessor, TokenData
from pretor.core.database.table.individual import AgentType
from fastapi import HTTPException
from typing import Optional, List, Dict
from pretor.utils.check_user.role_check import RoleChecker
from pretor.core.database.table.user import UserAuthority

agent_router = APIRouter(prefix="/api/v1/agent", tags=["agent"])

class AgentRegister(BaseModel):
    provider_title: str
    model_id: str
    individual_name: str

class AgentLocalRegister(BaseModel):
    path: str
    individual_name: str

@agent_router.post("")
async def load_agent(agent_register: Union[AgentRegister, AgentLocalRegister],
                     _: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER))):
    global_state_machine = ray_actor_hook("global_state_machine")
    if isinstance(agent_register, AgentLocalRegister):
        pass

    elif isinstance(agent_register, AgentRegister):
        match agent_register.individual_title:
            case "supervisory_node":
                node = ray_actor_hook("supervisory_node")
                node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
            case "consciousness_node":
                node = ray_actor_hook("consciousness_node")
                node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
            case "control_node":
                node = ray_actor_hook("control_node")
                node.create_agent.remote(global_state_machine,agent_register.provider_title,agent_register.model_id)
            case _:
                pass
    return {"message": "创建成功"}


class WorkerIndividualCreate(BaseModel):
    agent_name: str
    agent_type: AgentType
    description: str
    provider_title: str
    model_id: str
    system_prompt: str
    output_template: dict
    bound_skill: Dict[str, List[str]]
    workspace: List[str]


class WorkerIndividualUpdate(BaseModel):
    agent_name: Optional[str] = None
    agent_type: Optional[AgentType] = None
    description: Optional[str] = None
    provider_title: Optional[str] = None
    model_id: Optional[str] = None
    system_prompt: Optional[str] = None
    output_template: Optional[dict] = None
    bound_skill: Optional[Dict[str, List[str]]] = None
    workspace: Optional[List[str]] = None


@agent_router.post("/worker")
async def create_worker_individual(worker_data: WorkerIndividualCreate,
                                   token_data: TokenData = Depends(RoleChecker(allowed_roles=UserAuthority.USER))):
    postgres_database = ray_actor_hook("postgres_database")
    data_dict = worker_data.model_dump()
    data_dict["owner_id"] = token_data.user_id
    worker = await postgres_database.individual_database.remote("add_worker_individual", **data_dict)
    return {"message": "success", "agent_id": worker.agent_id}


@agent_router.get("/worker")
async def get_worker_individual_list(token_data: TokenData = Depends(Accessor.get_current_user)):
    postgres_database = ray_actor_hook("postgres_database")
    workers = await postgres_database.individual_database.remote("get_worker_individual_list", owner_id=token_data.user_id)
    return {"workers": workers}


@agent_router.get("/worker/{agent_id}")
async def get_worker_individual(agent_id: str,
                                token_data: TokenData = Depends(Accessor.get_current_user)):
    postgres_database = ray_actor_hook("postgres_database")
    worker = await postgres_database.individual_database.remote("get_worker_individual", agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")
    return worker


@agent_router.put("/worker/{agent_id}")
async def update_worker_individual(agent_id: str,
                                   worker_data: WorkerIndividualUpdate,
                                   token_data: TokenData = Depends(Accessor.get_current_user)):
    postgres_database = ray_actor_hook("postgres_database")
    worker = await postgres_database.individual_database.remote("get_worker_individual", agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")

    update_data = worker_data.model_dump(exclude_unset=True)
    updated_worker = await postgres_database.individual_database.remote("update_worker_individual", agent_id=agent_id, **update_data)
    return {"message": "success", "worker": updated_worker}


@agent_router.delete("/worker/{agent_id}")
async def delete_worker_individual(agent_id: str,
                                   token_data: TokenData = Depends(Accessor.get_current_user)):
    postgres_database = ray_actor_hook("postgres_database")
    worker = await postgres_database.individual_database.remote("get_worker_individual", agent_id=agent_id)
    if not worker:
        raise HTTPException(status_code=404, detail="Agent not found")
    if worker.owner_id != token_data.user_id:
        raise HTTPException(status_code=403, detail="Forbidden: You do not own this agent")
    await postgres_database.individual_database.remote("delete_worker_individual", agent_id=agent_id)
    return {"message": "success"}