"""User-facing importer/exporter types that translate to the API's connector wire format."""

from __future__ import annotations

import base64
import io
import tarfile
from pathlib import Path

from pydantic import BaseModel

from .types import (
    AgentWorkspaceDrainSpec,
    DataBallSourceSpec,
    ExportConnectorSpec,
    ImportConnectorSpec,
    RunSnapshotSourceSpec,
    S3DumpDrainSpec,
)


class LocalImporter(BaseModel):
    path: str

    def _to_wire(self) -> ImportConnectorSpec:
        encoded = _tarball_and_encode(self.path)
        return ImportConnectorSpec(
            source=DataBallSourceSpec(data=encoded),
            drain=AgentWorkspaceDrainSpec(),
        )


class S3Exporter(BaseModel):
    # when no bucket/path provided, the API uses a default bucket
    # FIXME: enable users to actually provide a bucket/path w/ bucket allocation, when that's implemented
    bucket: str | None = None
    path: str | None = None

    def _to_wire(self) -> ExportConnectorSpec:
        return ExportConnectorSpec(
            source=RunSnapshotSourceSpec(),
            drain=S3DumpDrainSpec(bucket=self.bucket, prefix=self.path),
        )


Importer = LocalImporter
Exporter = S3Exporter


def _tarball_and_encode(local_path: str) -> str:
    p = Path(local_path)
    buf = io.BytesIO()
    tf = tarfile.open(fileobj=buf, mode="w:gz")
    if p.is_dir():
        tf.add(str(p), arcname=".")
    else:
        tf.add(str(p), arcname=p.name)
    tf.close()
    return base64.b64encode(buf.getvalue()).decode("ascii")
