#!/usr/bin/env python3
"""Demo: seed a pool with LLM configs and prompts, fill it with real provider calls.

Usage:
  uv run python scripts/demo-pool-fill.py
  uv run python scripts/demo-pool-fill.py --dsn postgresql://postgres:postgres@localhost:5433/dr_llm_test

When no --dsn is provided, the script auto-creates a Docker-managed Postgres
project via ProjectInfo, runs the demo, and destroys it on exit.

The demo:
  - Defines LlmConfig instances for cheap/fast models
  - Defines short prompts as Message lists
  - Seeds the pending queue with the full (llm_config x prompt) cross product
  - Starts background workers using make_llm_process_fn (real LLM calls)
  - Prints progress until the queued work is complete
"""

from __future__ import annotations

import shutil
import time
from typing import Annotated
from uuid import uuid4

import typer

from dr_llm.pool import (
    DbConfig,
    DbRuntime,
    KeyColumn,
    PoolSchema,
    PoolStore,
    make_llm_process_fn,
    seed_pending,
    start_workers,
)
from dr_llm.pool.sample_models import WorkerSnapshot
from dr_llm.project.project_info import ProjectInfo
from dr_llm.providers import build_default_registry
from dr_llm.providers.llm_config import LlmConfig
from dr_llm.providers.models import Message

app = typer.Typer()

LLM_CONFIGS: dict[str, LlmConfig] = {
    "gpt-4.1-mini": LlmConfig(
        provider="openai",
        model="gpt-4.1-mini",
        temperature=0.7,
        max_tokens=64,
    ),
    "gemini-flash": LlmConfig(
        provider="google",
        model="gemini-2.5-flash",
        max_tokens=64,
    ),
}

PROMPTS: dict[str, list[Message]] = {
    "haiku": [Message(role="user", content="Write a haiku about programming.")],
    "math": [Message(role="user", content="What is 17 * 23? Reply with just the number.")],
}


def _print_progress(snapshot: WorkerSnapshot) -> None:
    counts = snapshot.status_counts
    print(
        "Progress: "
        f"claimed={snapshot.claimed} "
        f"promoted={snapshot.promoted} "
        f"failed={snapshot.failed} "
        f"pending={counts.pending} "
        f"leased={counts.leased}"
    )


def _run_demo(dsn: str, pool_name: str, num_workers: int, samples_per_cell: int) -> None:
    schema = PoolSchema(
        name=pool_name,
        key_columns=[KeyColumn(name="llm_config"), KeyColumn(name="prompt")],
    )
    runtime = DbRuntime(DbConfig(dsn=dsn))
    registry = build_default_registry()
    store = PoolStore(schema, runtime)
    controller = None

    try:
        store.init_schema()
        seed_result = seed_pending(
            store,
            key_grid={
                "llm_config": LLM_CONFIGS,
                "prompt": PROMPTS,
            },
            n=samples_per_cell,
        )
        print(
            f"Seeded {seed_result.inserted} pending rows"
            f" (skipped {seed_result.skipped} existing rows)"
        )

        process_fn = make_llm_process_fn(registry)
        controller = start_workers(
            store,
            process_fn=process_fn,
            num_workers=num_workers,
            min_poll_interval_s=0.5,
            max_poll_interval_s=3.0,
            max_retries=1,
        )
        try:
            last_progress: tuple[int, int, int, int, int] | None = None
            while True:
                snapshot = controller.snapshot()
                current_progress = (
                    snapshot.claimed,
                    snapshot.promoted,
                    snapshot.failed,
                    snapshot.status_counts.pending,
                    snapshot.status_counts.leased,
                )
                if current_progress != last_progress:
                    _print_progress(snapshot)
                    last_progress = current_progress
                if (
                    snapshot.status_counts.pending == 0
                    and snapshot.status_counts.leased == 0
                ):
                    break
                time.sleep(0.5)
        finally:
            controller.stop()
            final_snapshot = controller.join()

        final_counts = final_snapshot.status_counts
        print(
            "Final queue counts: "
            f"pending={final_counts.pending} "
            f"leased={final_counts.leased} "
            f"promoted={final_counts.promoted} "
            f"failed={final_counts.failed}"
        )

        coverage = store.coverage()
        total_samples = sum(row.count for row in coverage)
        print(f"Stored {total_samples} samples across {len(coverage)} cells")

        samples = store.bulk_load()
        for sample in samples[:4]:
            text = sample.payload.get("text", "")[:80]
            print(
                f"  [{sample.key_values['llm_config']}] "
                f"[{sample.key_values['prompt']}] "
                f"-> {text!r}"
            )
    finally:
        if controller is not None:
            controller.stop()
            controller.join()
        registry.close()
        runtime.close()


@app.command()
def main(
    dsn: Annotated[
        str | None,
        typer.Option(
            help="PostgreSQL DSN. If omitted, a temporary Docker project is created."
        ),
    ] = None,
    pool_name: Annotated[
        str,
        typer.Option(help="Pool name to create for the demo."),
    ] = "demo_pool_fill",
    num_workers: Annotated[
        int,
        typer.Option(help="Number of concurrent workers to run."),
    ] = 2,
    samples_per_cell: Annotated[
        int,
        typer.Option(help="Number of samples to queue for each (llm_config, prompt) cell."),
    ] = 1,
) -> None:
    if dsn is not None:
        _run_demo(dsn, pool_name, num_workers, samples_per_cell)
        return

    # Auto-manage a Docker Postgres project
    if not shutil.which("docker"):
        print("Error: Docker is required when no --dsn is provided.")
        print("Either install Docker or pass --dsn to use an existing database.")
        raise typer.Exit(1)

    project_name = f"demo-pool-fill-{uuid4().hex[:8]}"
    project: ProjectInfo | None = None
    try:
        print(f"Creating temporary project '{project_name}'...")
        project = ProjectInfo.create_new(project_name)
        assert project.dsn is not None
        print(f"Postgres ready at {project.dsn}")
        _run_demo(project.dsn, pool_name, num_workers, samples_per_cell)
    finally:
        if project is not None:
            print(f"Destroying temporary project '{project_name}'...")
            project.destroy()


if __name__ == "__main__":
    app()
