import asyncio
import podwrangler as pw

async def main():
    provider = pw.PlainTTYProvider()
    wss_url = "wss://k8s.provider.com/exec/45582/?token=<JWT>&cols=220&pod=my-app-abc-xyz"

    async with pw.PodSession.from_url(wss_url, provider) as session:
        result = await session.run("hostname")
        print(result.output)

asyncio.run(main())

provider = pw.PlainTTYProvider(cols=220)

provider = pw.KubernetesProvider(
    host="k8s.example.com:6443",
    use_tls=True,
    shell_cmd=["/bin/sh"],
)

async with pw.PodSession.from_url(wss_url, provider) as session:
    # Run a single command
    result = await session.run("ls -la /app")
    print(result.output)

    # Run with timeout and strict exit-code check
    result = await session.run("make test", timeout=120.0, check=True)

    # Run a sequence of commands
    results = await session.run_many(
        ["cd /app", "git pull", "pip install -r requirements.txt"],
        timeout_each=60.0,
        stop_on_error=True,
    )
    for r in results:
        print("✓" if r.succeeded else "✗", r.command)

registry = pw.PodRegistry()
registry.register_from_url("web-1", WSS_URL_1, provider)
registry.register_from_url("web-2", WSS_URL_2, provider)

print(registry.aliases)   # ['web-1', 'web-2']

async with registry.session("web-1") as s:
    result = await s.run("uptime")
    print(result.output)

async def run_on(alias, registry, cmd):
    async with registry.session(alias) as s:
        return await s.run(cmd)

results = await asyncio.gather(
    run_on("web-1", registry, "hostname"),
    run_on("web-2", registry, "hostname"),
)
for r in results:
    print(r.pod_label, "→", r.output.strip())

try:
    result = await session.run("ls /missing", check=True)
except pw.PodCommandTimeoutError as e:
    print(f"Timed out after {e.timeout}s")
except pw.PodCommandError as e:
    print(f"Exit {e.exit_code}: {e.output}")
except pw.PodConnectionError as e:
    print(f"Connection error: {e.reason}")

session = pw.PodSession(endpoint, auto_reconnect=True, connect_timeout=20.0)
await session.connect()
result = await session.run("echo hello")   # reconnects automatically if needed
await session.close()

provider = pw.PlainTTYProvider(cols=220)

identity = pw.PodIdentity(
    pod_name="my-app-deploy-abc12-xyz99",
    service="my-app-deploy",
    deployment="abc12",
    cluster="45582",   # PlainTTY service_id
    cols=220,
)

endpoint = pw.PodEndpoint(identity=identity, provider=provider, token="<JWT>")

async with pw.PodSession(endpoint, namespace="production") as s:
    r = await s.run("env | grep K8S")
    print(r.output)

class MyProvider:
    name = "myprovider"

    def build_wss_schema(self, identity: pw.PodIdentity, token: str) -> pw.WssSchema:
        url = f"wss://my.service.com/exec/{identity.pod_name}?token={token}"
        return pw.WssSchema(url=url, is_binary_framed=False)

    def parse_url(self, url: str) -> tuple[pw.PodIdentity, str]:
        ...

    async def init_shell(self, ws, schema: pw.WssSchema) -> None:
        ...

    async def detect_namespace(self, ws, identity, schema) -> str:
        return identity.namespace or "default"

    async def list_pods(self, token: str, **kwargs) -> list[pw.PodIdentity]:
        raise NotImplementedError

async with pw.PodSession.from_url(my_url, MyProvider()) as s:
    print(await s.run("hostname"))