from fastapi import APIRouter
from fastapi import HTTPException, Depends ,Security

from fastapi_sqlalchemy import db
from pydantic import BaseModel

from superagi.worker import execute_agent
from superagi.helper.auth import validate_api_key,get_organisation_from_api_key
from superagi.models.agent import Agent
from superagi.models.agent_execution_config import AgentExecutionConfiguration
from superagi.models.agent_config import AgentConfiguration
from superagi.models.agent_schedule import AgentSchedule
from superagi.models.project import Project
from superagi.models.workflows.agent_workflow import AgentWorkflow
from superagi.models.agent_execution import AgentExecution
from superagi.models.organisation import Organisation
from superagi.models.knowledges import Knowledges
from superagi.models.resource import Resource
from superagi.controllers.types.agent_with_config import AgentConfigExtInput,AgentConfigUpdateExtInput
from superagi.models.workflows.iteration_workflow import IterationWorkflow
from superagi.helper.s3_helper import S3Helper
from datetime import datetime
from typing import Optional,List
from superagi.models.toolkit import Toolkit
from superagi.apm.event_handler import EventHandler
from superagi.config.config import get_config
router = APIRouter()

class AgentExecutionIn(BaseModel):
    name: Optional[str]
    goal: Optional[List[str]]
    instruction: Optional[List[str]]

    class Config:
        orm_mode = True

class RunFilterConfigIn(BaseModel):
    run_ids:Optional[List[int]]
    run_status_filter:Optional[str]

    class Config:
        orm_mode = True

class ExecutionStateChangeConfigIn(BaseModel):
    run_ids:Optional[List[int]]

    class Config:
        orm_mode = True

class RunIDConfig(BaseModel):
    run_ids:List[int]

    class Config:
        orm_mode = True

@router.post("", status_code=200)
def create_agent_with_config(agent_with_config: AgentConfigExtInput,
                             api_key: str = Security(validate_api_key), organisation:Organisation = Depends(get_organisation_from_api_key)):
    project=Project.find_by_org_id(db.session, organisation.id)
    try:
        tools_arr=Toolkit.get_tool_and_toolkit_arr(db.session,organisation.id,agent_with_config.tools)
    except Exception as e:
        raise HTTPException(status_code=404, detail=str(e))

    agent_with_config.tools=tools_arr
    agent_with_config.project_id=project.id
    agent_with_config.exit="No exit criterion"
    agent_with_config.permission_type="God Mode"
    agent_with_config.LTM_DB=None
    db_agent = Agent.create_agent_with_config(db, agent_with_config)

    if agent_with_config.schedule is not None:
        agent_schedule = AgentSchedule.save_schedule_from_config(db.session, db_agent, agent_with_config.schedule)
        if agent_schedule is None:
            raise HTTPException(status_code=500, detail="Failed to schedule agent")
        EventHandler(session=db.session).create_event('agent_created', {'agent_name': agent_with_config.name,
                                                                            'model': agent_with_config.model}, db_agent.id,
                                                        organisation.id if organisation else 0)
        db.session.commit()
        return {
            "agent_id": db_agent.id
        }
    
    start_step = AgentWorkflow.fetch_trigger_step_id(db.session, db_agent.agent_workflow_id)
    iteration_step_id = IterationWorkflow.fetch_trigger_step_id(db.session,
                                                                start_step.action_reference_id).id if start_step.action_type == "ITERATION_WORKFLOW" else -1
    # Creating an execution with RUNNING status
    execution = AgentExecution(status='CREATED', last_execution_time=datetime.now(), agent_id=db_agent.id,
                               name="New Run", current_agent_step_id=start_step.id, iteration_workflow_step_id=iteration_step_id)
    agent_execution_configs = {
        "goal": agent_with_config.goal,
        "instruction": agent_with_config.instruction,
        "constraints": agent_with_config.constraints,
        "exit": agent_with_config.exit,
        "tools": agent_with_config.tools,
        "iteration_interval": agent_with_config.iteration_interval,
        "model": agent_with_config.model,
        "permission_type": agent_with_config.permission_type,
        "LTM_DB": agent_with_config.LTM_DB,
        "max_iterations": agent_with_config.max_iterations,
        "user_timezone": agent_with_config.user_timezone,
        "knowledge": agent_with_config.knowledge
    }
    db.session.add(execution)
    db.session.commit()
    db.session.flush()
    AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=execution,
                                                                     agent_execution_configs=agent_execution_configs)

    organisation = db_agent.get_agent_organisation(db.session)
    EventHandler(session=db.session).create_event('agent_created', {'agent_name': agent_with_config.name,
                                                                    'model': agent_with_config.model}, db_agent.id,
                                                  organisation.id if organisation else 0)
    # execute_agent.delay(execution.id, datetime.now())
    db.session.commit()
    return {
        "agent_id": db_agent.id
    }

@router.post("/{agent_id}/run",status_code=200)
def create_run(agent_id:int,agent_execution: AgentExecutionIn,api_key: str = Security(validate_api_key),organisation:Organisation = Depends(get_organisation_from_api_key)):
    agent=Agent.get_agent_from_id(db.session,agent_id)
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    project=Project.find_by_id(db.session, agent.project_id)
    if project.organisation_id!=organisation.id:
        raise HTTPException(status_code=404, detail="Agent not found")
    db_schedule=AgentSchedule.find_by_agent_id(db.session, agent_id)
    if db_schedule is not None:
        raise HTTPException(status_code=409, detail="Agent is already scheduled,cannot run")
    start_step = AgentWorkflow.fetch_trigger_step_id(db.session, agent.agent_workflow_id)
    db_agent_execution=AgentExecution.get_execution_by_agent_id_and_status(db.session, agent_id, "CREATED")

    if db_agent_execution is None:
        db_agent_execution = AgentExecution(status="RUNNING", last_execution_time=datetime.now(),
                                            agent_id=agent_id, name=agent_execution.name, num_of_calls=0,
                                            num_of_tokens=0,
                                            current_agent_step_id=start_step.id)
        db.session.add(db_agent_execution)
    else:
        db_agent_execution.status = "RUNNING"

    db.session.commit()
    db.session.flush()

    agent_execution_configs = {}
    if agent_execution.goal is not None:
        agent_execution_configs = {
            "goal": agent_execution.goal,
        }

    if agent_execution.instruction is not None:
        agent_execution_configs["instructions"] = agent_execution.instruction,

    if agent_execution_configs != {}:
        AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=db_agent_execution,
                                                                     agent_execution_configs=agent_execution_configs)
    EventHandler(session=db.session).create_event('run_created', 
                                                  {'agent_execution_id': db_agent_execution.id,
                                                   'agent_execution_name':db_agent_execution.name
                                                   },
                                                   agent_id, 
                                                   organisation.id if organisation else 0)
    
    agent_execution_knowledge = AgentConfiguration.get_agent_config_by_key_and_agent_id(session= db.session, key= 'knowledge', agent_id= agent_id)
    if agent_execution_knowledge and agent_execution_knowledge.value != 'None':
        knowledge_name = Knowledges.get_knowledge_from_id(db.session, int(agent_execution_knowledge.value)).name
        if knowledge_name is not None:
            EventHandler(session=db.session).create_event('knowledge_picked', 
                                                        {'knowledge_name': knowledge_name, 
                                                         'agent_execution_id': db_agent_execution.id},
                                                        agent_id,
                                                        organisation.id if organisation else 0
                                                        )

    if db_agent_execution.status == "RUNNING":
      execute_agent.delay(db_agent_execution.id, datetime.now())
    return {
        "run_id":db_agent_execution.id
    }

@router.put("/{agent_id}",status_code=200)
def update_agent(agent_id: int, agent_with_config: AgentConfigUpdateExtInput,api_key: str = Security(validate_api_key),
                                        organisation:Organisation = Depends(get_organisation_from_api_key)):
    
    db_agent= Agent.get_active_agent_by_id(db.session, agent_id)
    if not db_agent:
        raise HTTPException(status_code=404, detail="agent not found")
    
    project=Project.find_by_id(db.session, db_agent.project_id)
    if project is None:
        raise HTTPException(status_code=404, detail="Project not found")
    
    if project.organisation_id!=organisation.id:
        raise HTTPException(status_code=404, detail="Agent not found")

    # db_execution=AgentExecution.get_execution_by_agent_id_and_status(db.session, agent_id, "RUNNING")
    # if db_execution is not None:
    #     raise HTTPException(status_code=409, detail="Agent is already running,please pause and then update")
     
    db_schedule=AgentSchedule.find_by_agent_id(db.session, agent_id)
    if db_schedule is not None:
        raise HTTPException(status_code=409, detail="Agent is already scheduled,cannot update")
    
    try:
        tools_arr=Toolkit.get_tool_and_toolkit_arr(db.session,organisation.id,agent_with_config.tools)
    except Exception as e:
        raise HTTPException(status_code=404,detail=str(e))

    if agent_with_config.schedule is not None:
        raise HTTPException(status_code=400,detail="Cannot schedule an existing agent")
    agent_with_config.tools=tools_arr
    agent_with_config.project_id=project.id
    agent_with_config.exit="No exit criterion"
    agent_with_config.permission_type="God Mode"
    agent_with_config.LTM_DB=None

    for key,value in agent_with_config.dict().items():
        if hasattr(db_agent,key) and value is not None:
            setattr(db_agent,key,value)
    db.session.commit()
    db.session.flush()

    start_step = AgentWorkflow.fetch_trigger_step_id(db.session, db_agent.agent_workflow_id)
    iteration_step_id = IterationWorkflow.fetch_trigger_step_id(db.session,
                                                                start_step.action_reference_id).id if start_step.action_type == "ITERATION_WORKFLOW" else -1
    execution = AgentExecution(status='CREATED', last_execution_time=datetime.now(), agent_id=db_agent.id,
                               name="New Run", current_agent_step_id=start_step.id, iteration_workflow_step_id=iteration_step_id)
    agent_execution_configs = {
        "goal": agent_with_config.goal,
        "instruction": agent_with_config.instruction,
        "tools":agent_with_config.tools,
        "constraints": agent_with_config.constraints,
        "iteration_interval": agent_with_config.iteration_interval,
        "model": agent_with_config.model,
        "max_iterations": agent_with_config.max_iterations,
        "agent_workflow": agent_with_config.agent_workflow,
    }
    agent_configurations = [
        AgentConfiguration(agent_id=db_agent.id, key=key, value=str(value))
        for key, value in agent_execution_configs.items()
    ]
    db.session.add_all(agent_configurations)
    db.session.add(execution)
    db.session.commit()
    db.session.flush()
    AgentExecutionConfiguration.add_or_update_agent_execution_config(session=db.session, execution=execution,
                                                                     agent_execution_configs=agent_execution_configs)
    db.session.commit()

    return {
        "agent_id":db_agent.id
    }


@router.post("/{agent_id}/run-status")
def get_agent_runs(agent_id:int,filter_config:RunFilterConfigIn,api_key: str = Security(validate_api_key),organisation:Organisation = Depends(get_organisation_from_api_key)):
    agent= Agent.get_active_agent_by_id(db.session, agent_id)
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    project=Project.find_by_id(db.session, agent.project_id)
    if project.organisation_id!=organisation.id:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    db_execution_arr=[]
    if filter_config.run_status_filter is not None:
        filter_config.run_status_filter=filter_config.run_status_filter.upper()

    db_execution_arr=AgentExecution.get_all_executions_by_filter_config(db.session, agent.id, filter_config)
    
    response_arr=[]
    for ind_execution in db_execution_arr:
        response_arr.append({"run_id":ind_execution.id, "status":ind_execution.status})

    return response_arr


@router.post("/{agent_id}/pause",status_code=200)
def pause_agent_runs(agent_id:int,execution_state_change_input:ExecutionStateChangeConfigIn,api_key: str = Security(validate_api_key),organisation:Organisation = Depends(get_organisation_from_api_key)):
    agent= Agent.get_active_agent_by_id(db.session, agent_id)
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    project=Project.find_by_id(db.session, agent.project_id)
    if project.organisation_id!=organisation.id:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    #Checking if the run_ids whose output files are requested belong to the organisation 
    if execution_state_change_input.run_ids is not None:
        try:
            AgentExecution.validate_run_ids(db.session,execution_state_change_input.run_ids,organisation.id)
        except Exception as e:
            raise HTTPException(status_code=404, detail="One or more run id(s) not found")
    
    db_execution_arr=AgentExecution.get_all_executions_by_status_and_agent_id(db.session, agent.id, execution_state_change_input, "RUNNING")

    if db_execution_arr is not None and execution_state_change_input.run_ids is not None \
            and len(db_execution_arr) != len(execution_state_change_input.run_ids):
        raise HTTPException(status_code=404, detail="One or more run id(s) not found")

    for ind_execution in db_execution_arr:
        ind_execution.status="PAUSED"
    db.session.commit()
    db.session.flush()
    return {
        "result":"success"
    }

@router.post("/{agent_id}/resume",status_code=200)
def resume_agent_runs(agent_id:int,execution_state_change_input:ExecutionStateChangeConfigIn,api_key: str = Security(validate_api_key),organisation:Organisation = Depends(get_organisation_from_api_key)):
    agent= Agent.get_active_agent_by_id(db.session, agent_id)
    if not agent:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    project=Project.find_by_id(db.session, agent.project_id)
    if project.organisation_id!=organisation.id:
        raise HTTPException(status_code=404, detail="Agent not found")
    
    if execution_state_change_input.run_ids is not None:
        try:
            AgentExecution.validate_run_ids(db.session,execution_state_change_input.run_ids,organisation.id)
        except Exception as e:
            raise HTTPException(status_code=404, detail="One or more run id(s) not found")
    
    db_execution_arr=AgentExecution.get_all_executions_by_status_and_agent_id(db.session, agent.id, execution_state_change_input, "PAUSED")

    if db_execution_arr is not None and execution_state_change_input.run_ids is not None\
            and len(db_execution_arr) != len(execution_state_change_input.run_ids):
        raise HTTPException(status_code=404, detail="One or more run id(s) not found")

    for ind_execution in db_execution_arr:
        ind_execution.status="RUNNING"
        execute_agent.delay(ind_execution.id, datetime.now())
        
    db.session.commit()
    db.session.flush()

    return {
        "result":"success"
    }

@router.post("/resources/output",status_code=200)
def get_run_resources(run_id_config:RunIDConfig,api_key: str = Security(validate_api_key),organisation:Organisation = Depends(get_organisation_from_api_key)):
    if get_config('STORAGE_TYPE') != "S3":
        raise HTTPException(status_code=400,detail="This endpoint only works when S3 is configured")
    run_ids_arr=run_id_config.run_ids
    if len(run_ids_arr)==0:  
        raise HTTPException(status_code=404,
                            detail=f"No execution_id found")
    #Checking if the run_ids whose output files are requested belong to the organisation 
    try:
        AgentExecution.validate_run_ids(db.session, run_ids_arr, organisation.id)
    except Exception as e:
        raise HTTPException(status_code=404, detail="One or more run id(s) not found")
    
    db_resources_arr=Resource.find_by_run_ids(db.session, run_ids_arr)

    try:
        response_obj=S3Helper().get_download_url_of_resources(db_resources_arr)
    except:
        raise HTTPException(status_code=401,detail="Invalid S3 credentials")
    return response_obj