#  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

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(Accessor.get_current_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(Accessor.get_current_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.add_worker_individual.remote(**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.get_worker_individual_list.remote(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.get_worker_individual.remote(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.get_worker_individual.remote(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.update_worker_individual.remote(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.get_worker_individual.remote(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.delete_worker_individual.remote(agent_id=agent_id)
    return {"message": "success"}