import asyncio
import httpx
import requests
import datetime
import json
import uuid
import time
from typing import List
import multiprocessing
from sqlmodel import Session, select
from a2a.client import A2ACardResolver, A2AClient
from a2a.types import (
    AgentCard,
    Part,
    TextPart,
    MessageSendParams,
    Message,
    Role,
    SendMessageRequest,
    SendMessageResponse,
)

from agentbeats.model import (
    Agent,
    AgentCheck,
    engine,
    Assessment,
    AssessmentStatus,
    LogEntry,
)
from agentbeats.settings import ASSESSMENT_RUNNER_TIMEOUT


def check_agent(session: Session, agent: Agent):
    print(f"Creating AgentCheck for Agent {agent.id}...")
    is_ctrl_reachable = False
    agent_count = None
    agent_url = None
    card_url = None
    card_content = None

    # check if ctrl_url is reachable
    try:
        assert agent.ctrl_url is not None
        ctrl_status_url = f"{agent.ctrl_url}/status"
        resp = requests.get(ctrl_status_url, timeout=5)
        if resp.status_code == 200:
            is_ctrl_reachable = True
            data = resp.json()
            agent_count = data.get("running_agents", None)
            if agent_count and agent_count >= 1:
                agents_url = f"{agent.ctrl_url}/agents"
                resp_agents = requests.get(agents_url, timeout=5)
                if resp_agents.status_code == 200:
                    agents_data = resp_agents.json()
                    if len(agents_data) > 0:
                        # FIXME: support scaling later
                        agent_info = list(agents_data.values())[0]
                        state = agent_info.get("state", None)
                        if state != "running":
                            print(f"Agent {agent.id} state is not running: {state}")
                            raise Exception("Agent not running")
                        agent_url = agent_info.get("url", None)
                        card_url = f"{agent_url}/.well-known/agent-card.json"
                        resp_card = requests.get(card_url, timeout=5)
                        if resp_card.status_code == 200:
                            card_content = json.dumps(resp_card.json())
    except Exception as e:
        import traceback

        traceback.print_exc()
        print(f"Error checking agent {agent.id}: {e}")

    assert agent.id is not None
    check = AgentCheck(
        agent_id=agent.id,
        created_at=datetime.datetime.utcnow(),
        is_ctrl_reachable=is_ctrl_reachable,
        agent_count=agent_count,
        agent_url=agent_url,
        card_url=card_url,
        card_content=card_content,
    )
    session.add(check)
    session.commit()
    print(f"Created AgentCheck for Agent {agent.id}")


def log(session, assessment_id, source, channel, content):
    log_entry = LogEntry(
        assessment_id=assessment_id,
        source=source,
        channel=channel,
        content=content,
        created_at=datetime.datetime.utcnow(),
    )
    session.add(log_entry)
    session.commit()
    return log_entry


def get_first_cagent_id(ctrl_url: str):
    agents_url = f"{ctrl_url}/agents"
    resp = requests.get(agents_url, timeout=5)
    if resp.status_code == 200:
        agents_data = resp.json()
        if len(agents_data) > 0:
            return list(agents_data.keys())[0]
    return None


def reset_cagent(ctrl_url: str, cagent_id: str):
    reset_url = f"{ctrl_url}/agents/{cagent_id}/reset"
    resp = requests.post(reset_url, timeout=5)
    if resp.status_code == 200:
        return True
    return False


def sync_clog_log(target_url: str, log_ids: List[uuid.UUID], keys: List[str]):
    with Session(engine) as session:
        try:
            resp = requests.get(target_url, timeout=5)
            if resp.status_code == 200:
                for log_entry_id, key in zip(log_ids, keys):
                    log_entry = session.get(LogEntry, log_entry_id)
                    assert log_entry is not None
                    clog = resp.json().get(key, "")
                    if clog != log_entry.content:
                        log_entry.content = clog
                        session.add(log_entry)
                        session.commit()
        except Exception as e:
            pass


def sync_clog_log_loop(target_url: str, log_ids: List[uuid.UUID], keys: List[str]):
    while True:
        time.sleep(1)
        sync_clog_log(target_url, log_ids, keys)


async def send_a2a_message(
    url, message, task_id=None, context_id=None
) -> SendMessageResponse:
    httpx_client = httpx.AsyncClient(timeout=1800.0)
    resolver = A2ACardResolver(httpx_client=httpx_client, base_url=url)
    card: AgentCard | None = await resolver.get_agent_card()
    client = A2AClient(httpx_client=httpx_client, agent_card=card)
    message_id = uuid.uuid4().hex
    params = MessageSendParams(
        message=Message(
            role=Role.user,
            parts=[Part(TextPart(text=message))],
            message_id=message_id,
            task_id=task_id,
            context_id=context_id,
        )
    )
    request_id = uuid.uuid4().hex
    req = SendMessageRequest(id=request_id, params=params)
    response = await client.send_message(request=req)
    return response


def kickoff(session: Session, assessment: Assessment, cagent_info: list):
    sync_ps = []
    # start log sync
    for ag, ctrl_url, cagent_id in cagent_info:
        stdout_log = log(session, assessment.id, str(ag.id), "stdout", "")
        stderr_log = log(session, assessment.id, str(ag.id), "stderr", "")
        target_url = f"{ctrl_url}/agents/{cagent_id}"
        p = multiprocessing.Process(
            target=sync_clog_log_loop,
            args=(
                target_url,
                [stdout_log.id, stderr_log.id],
                ["stdout_log", "stderr_log"],
            ),
        )
        p.start()
        sync_ps.append(p)
    # prepare cagent_urls
    cagent_urls = []
    for ag, ctrl_url, cagent_id in cagent_info:
        cagent_url = f"{ctrl_url}/to_agent/{cagent_id}"
        cagent_urls.append(cagent_url)

    if assessment.config == "default":
        task_text = "Your task is to assess the agents located at:\n"
        for url in cagent_urls[1:]:
            task_text += f"<white_agent_url>\n{url}\n</white_agent_url>\n"

    elif assessment.config == "tau-bench":
        # FIXME: temporary hack, config should be editable from UI later
        task_config = {
            "env": "retail",
            "user_strategy": "llm",
            "user_model": "openrouter/openai/gpt-4o",
            "user_provider": "litellm_proxy",
            "task_split": "test",
            "task_ids": [1],
        }
        task_text = f"""
Your task is to instantiate tau-bench to test the agent located at:
<white_agent_url>
{cagent_urls[1]}
</white_agent_url>
You should use the following env configuration:
<env_config>
{json.dumps(task_config, indent=2)}
</env_config>"""
    else:
        raise Exception(f"Unknown assessment config: {assessment.config}")

    log(session, assessment.id, "platform", "task_start", task_text)
    result = asyncio.run(
        send_a2a_message(
            url=cagent_urls[0],
            message=task_text,
        )
    )
    log(session, assessment.id, "platform", "task_return", str(result))
    # wait for at least one round of sync
    time.sleep(2)
    # kill sync processes
    for p in sync_ps:
        p.terminate()


# need to make sure this terminates within 10 minutes
def run_assessment(assessment_id: str):
    p = multiprocessing.Process(target=_run_assessment, args=(assessment_id,))
    p.start()
    p.join(timeout=ASSESSMENT_RUNNER_TIMEOUT)
    if p.is_alive():
        p.terminate()
        p.join()
        try:
            with Session(engine) as session:
                assessment = session.get(Assessment, uuid.UUID(assessment_id))
                assert assessment is not None
                log(
                    session,
                    assessment.id,
                    "platform",
                    "runner_log",
                    "ASSESSMENT_RUNNER_TIMEOUT reached, terminated.",
                )
                if assessment is not None:
                    assessment.status = AssessmentStatus.ERROR
                    session.add(assessment)
                    session.commit()
        except Exception as e:
            import traceback

            traceback.print_exc()
            print(f"Error updating assessment status after timeout: {e}")


def _run_assessment(assessment_id: str):
    with Session(engine) as session:
        assessment = session.get(Assessment, uuid.UUID(assessment_id))
        assert assessment is not None

        def log_runner(msg):
            return log(session, assessment.id, "platform", "runner_log", msg)

        try:
            # Gather agent instances inside controllers
            log_runner(
                f"Starting assessment, gathering {len(assessment.agents)} agent instances..."
            )
            cagent_info = []
            for ag in assessment.agents:
                ctrl_url = ag.ctrl_url
                assert ctrl_url is not None
                cagent_id = get_first_cagent_id(ctrl_url)
                assert cagent_id is not None
                cagent_info.append((ag, ctrl_url, cagent_id))
                log_runner(
                    f"Found agent {ag.id} at controller {ctrl_url} with cagent_id {cagent_id}"
                )
            # reset all agents
            log_runner("Resetting all agents...")
            for ag, ctrl_url, cagent_id in cagent_info:
                success = reset_cagent(ctrl_url, cagent_id)
                if success:
                    log_runner(f"Reset agent {ag.id} successfully.")
                else:
                    log_runner(f"Failed to reset agent {ag.id}.")
                    raise Exception(f"Failed to reset agent {ag.id}")
            # wait until they are ready
            log_runner("All agents reset successfully, waiting for them to be ready...")
            for ag, ctrl_url, cagent_id in cagent_info:
                # check for 30 times, every 2 seconds
                ready = False
                for _ in range(30):
                    cagent_url = f"{ctrl_url}/agents/{cagent_id}"
                    resp = requests.get(cagent_url, timeout=5)
                    if resp.status_code == 200:
                        status_data = resp.json()
                        state = status_data.get("state", None)
                        if state == "running":
                            ready = True
                            log_runner(f"Agent {ag.id} is ready.")
                            break
                    time.sleep(2)
                if not ready:
                    log_runner(f"Agent {ag.id} is not ready after waiting.")
                    raise Exception(f"Agent {ag.id} is not ready")
            # send to green and wait for feedback (and add to log)
            log_runner("All agents are ready, running assessment...")
            assessment.status = AssessmentStatus.RUNNING
            session.add(assessment)
            session.commit()
            session.refresh(assessment)
            kickoff(session, assessment, cagent_info)
            # mark it finished
            assessment.status = AssessmentStatus.COMPLETED
            session.add(assessment)
            session.commit()
            log_runner("Assessment finished successfully.")
        except Exception as e:
            import traceback

            log_runner(f"Error running assessment: {e}")
            log_runner(f"Traceback: {traceback.format_exc()}")
            assessment.status = AssessmentStatus.ERROR
            session.add(assessment)
            session.commit()
