import logging
import time
from pathlib import Path
from typing import AsyncIterator, Optional

from agent_protocol_client import (
    AgentApi,
    ApiClient,
    Configuration,
    Step,
    TaskRequestBody,
)

from agbenchmark.agent_interface import get_list_of_file_paths
from agbenchmark.config import AgentBenchmarkConfig

logger = logging.getLogger(__name__)


async def run_api_agent(
    task: str,
    config: AgentBenchmarkConfig,
    timeout: int,
    artifacts_location: Optional[Path] = None,
    *,
    mock: bool = False,
) -> AsyncIterator[Step]:
    configuration = Configuration(host=config.host)
    async with ApiClient(configuration) as api_client:
        api_instance = AgentApi(api_client)
        task_request_body = TaskRequestBody(input=task, additional_input=None)

        start_time = time.time()
        response = await api_instance.create_agent_task(
            task_request_body=task_request_body
        )
        task_id = response.task_id

        if artifacts_location:
            logger.debug("Uploading task input artifacts to agent...")
            await upload_artifacts(
                api_instance, artifacts_location, task_id, "artifacts_in"
            )

        logger.debug("Running agent until finished or timeout...")
        while True:
            step = await api_instance.execute_agent_task_step(task_id=task_id)
            yield step

            if time.time() - start_time > timeout:
                raise TimeoutError("Time limit exceeded")
            if step and mock:
                step.is_last = True
            if not step or step.is_last:
                break

        if artifacts_location:
            # In "mock" mode, we cheat by giving the correct artifacts to pass the test
            if mock:
                logger.debug("Uploading mock artifacts to agent...")
                await upload_artifacts(
                    api_instance, artifacts_location, task_id, "artifacts_out"
                )

            logger.debug("Downloading agent artifacts...")
            await download_agent_artifacts_into_folder(
                api_instance, task_id, config.temp_folder
            )


async def download_agent_artifacts_into_folder(
    api_instance: AgentApi, task_id: str, folder: Path
):
    artifacts = await api_instance.list_agent_task_artifacts(task_id=task_id)

    for artifact in artifacts.artifacts:
        # current absolute path of the directory of the file
        if artifact.relative_path:
            path: str = (
                artifact.relative_path
                if not artifact.relative_path.startswith("/")
                else artifact.relative_path[1:]
            )
            folder = (folder / path).parent

        if not folder.exists():
            folder.mkdir(parents=True)

        file_path = folder / artifact.file_name
        logger.debug(f"Downloading agent artifact {artifact.file_name} to {folder}")
        with open(file_path, "wb") as f:
            content = await api_instance.download_agent_task_artifact(
                task_id=task_id, artifact_id=artifact.artifact_id
            )

            f.write(content)


async def upload_artifacts(
    api_instance: AgentApi, artifacts_location: Path, task_id: str, type: str
) -> None:
    for file_path in get_list_of_file_paths(artifacts_location, type):
        relative_path: Optional[str] = "/".join(
            str(file_path).split(f"{type}/", 1)[-1].split("/")[:-1]
        )
        if not relative_path:
            relative_path = None

        await api_instance.upload_agent_task_artifacts(
            task_id=task_id, file=str(file_path), relative_path=relative_path
        )