import time
import json
import requests
import multiprocessing


from sqlmodel import Session, select
from agentbeats.model import (
    engine,
    Agent,
    AgentHostedStatus,
    Assessment,
    AssessmentStatus,
)
from agentbeats.gcp import upload_docker_to_gcp, launch_in_cloud_run, GCPSettings
from agentbeats.myutils import check_agent, run_assessment


def cloudrun_runner_loop():
    # start runner loop after 5s (wait for the database to be ready)
    time.sleep(5)
    while True:
        time.sleep(1)
        with Session(engine) as session:
            # check the database for pending agents to launch
            statement = select(Agent).where(
                Agent.deploy_type == "hosted",
                Agent.hosted_status == "pending",
            )
            results = session.exec(statement)
            pending_agents = results.all()
            for agent in pending_agents:
                print(f"Found pending agent to launch: {agent.id}")
                try:
                    # upload docker image to GCP
                    gcp_image_url = upload_docker_to_gcp(agent)
                    print(f"Uploaded Docker image to GCP: {gcp_image_url}")
                    # launch in Cloud Run
                    inj_env = {}
                    if agent.secret:
                        inj_env = json.loads(agent.secret)
                    if agent.inject_litellm_proxy_api:
                        inj_env["LITELLM_PROXY_API_KEY"] = (
                            "sk-1drUa7BXyAibGKVo6PJYxQ"  # FIXME: use real key management
                        )
                        inj_env["LITELLM_PROXY_API_BASE"] = (
                            "https://dev1.agentbeats.org:12333"
                        )
                    response = launch_in_cloud_run(
                        agent,
                        gcp_image_url,
                        additional_env=inj_env,
                    )
                    print(f"Launched Agent {agent.id} in Cloud Run: {response}")
                    # update agent status in database
                    agent.hosted_status = AgentHostedStatus.STARTING
                    agent.gcp_docker_image_url = gcp_image_url
                    agent.gcp_cloudrun_service_name = "agent-" + str(agent.id)
                    gcp_settings = GCPSettings()
                    agent.ctrl_url = (
                        f"https://agent-{agent.id}-{gcp_settings.gcp_endpoint_suffix}"
                    )
                    session.add(agent)
                except Exception as e:
                    print(f"Error launching Agent {agent.id}: {e}")
                    agent.hosted_status = AgentHostedStatus.ERROR
                    agent.hosted_error_msg = str(e)
                    session.add(agent)
                session.commit()
                session.refresh(agent)
            # check the database for started agents to check ready
            statement = select(Agent).where(
                Agent.deploy_type == "hosted",
                Agent.hosted_status == AgentHostedStatus.STARTING,
            )
            results = session.exec(statement)
            starting_agents = results.all()
            for agent in starting_agents:
                print(f"Checking status of starting agent: {agent.id}")
                status_url = f"{agent.ctrl_url}/status"
                try:
                    resp = requests.get(status_url, timeout=5)
                    if resp.status_code == 200:
                        status_data = resp.json()
                        if status_data.get("running_agents") >= 1:
                            print(f"Agent {agent.id} is ready")
                            agent.hosted_status = AgentHostedStatus.RUNNING
                            session.add(agent)
                            session.commit()
                            session.refresh(agent)
                except Exception as _:
                    pass  # Not ready yet


def agent_check_runner_loop():
    # start runner loop after 5s (wait for the database to be ready)
    time.sleep(1)
    while True:
        time.sleep(1)
        try:
            # check if any agent does not have a binded agent check
            with Session(engine) as session:
                for agent in session.exec(select(Agent)).all():
                    if len(agent.agent_checks) == 0:
                        print(f"Found agent {agent.id} with no check.")
                        check_agent(session, agent)
        except Exception as e:
            print(f"Error in agent check runner loop: {e}")


def assessment_runner_loop(pool_size=4):
    # start runner loop after 5s (wait for the database to be ready)
    time.sleep(1)
    while True:
        time.sleep(1)
        with Session(engine) as session:
            # check for pending assessments
            selected_agents = set()
            assessments = session.exec(
                select(Assessment).where(Assessment.status == AssessmentStatus.PENDING)
            ).all()
            # filter out the ones that are not owned by account with berkeley.edu email
            todos = []
            for a in assessments:
                user = a.user
                emails = [
                    e["email"]
                    for e in json.loads(user.github_emails_json)
                    if e["verified"]
                ]
                if not any(
                    (
                        email.endswith("@berkeley.edu")
                        or email.endswith("@cs.berkeley.edu")
                    )
                    for email in emails
                ):
                    continue
                if any(ag.id in selected_agents for ag in a.agents):
                    continue  # only select non-conflicting assessments
                print(f"Found qualified pending assessment {a.id}, launching...")
                a.status = AssessmentStatus.STARTING
                session.add(a)
                session.commit()
                session.refresh(a)
                todos.append(str(a.id))
                for ag in a.agents:
                    selected_agents.add(ag.id)
            if len(todos) > 0:
                print(
                    f"Launching {len(todos)} assessments in pool of size {pool_size}: {todos}"
                )
                # but considering daemon process limitation, do not directly use pool
                procs = [
                    multiprocessing.Process(target=run_assessment, args=(aid,))
                    for aid in todos
                ]
                for p in procs:
                    p.start()
                for p in procs:
                    p.join()
