"""One-shot suggestion: read project state, build prompt, call LLM, submit."""

from __future__ import annotations

from .client import NbenchClient


def _build_prompt(project: dict, experiments: list, categories: list, code: str) -> str:
    """Assemble LLM prompt from project context."""
    lines = [
        f"You are an AI research assistant helping improve a machine-learning training script.",
        f"",
        f"## Project",
        f"Name: {project.get('title', 'unknown')}",
        f"Goal: {project.get('goal', 'Minimize val_bpb')}",
        f"",
    ]

    if categories:
        lines.append("## Experiment Categories (best val_bpb per category)")
        for cat in categories[:15]:
            bpb = cat.get("best_val_bpb", "?")
            lines.append(f"- {cat.get('name', '?')}: best {bpb}")
        lines.append("")

    if experiments:
        lines.append("## Recent Experiments")
        for exp in experiments[:20]:
            status = exp.get("status", "?")
            bpb = exp.get("val_bpb", "?")
            desc = exp.get("description", "")[:100]
            lines.append(f"- #{exp.get('experiment_number','?')} [{status}] bpb={bpb} {desc}")
        lines.append("")

    if code:
        lines.append("## Current train.py")
        lines.append("```python")
        lines.append(code[:8000])
        lines.append("```")
        lines.append("")

    lines.append(
        "Based on the above context, suggest ONE concrete code change to improve "
        "validation performance (lower val_bpb). Be specific — describe the exact "
        "modification and why it should help. Keep your response concise (under 500 words)."
    )
    return "\n".join(lines)


def _get_provider(provider_name: str, model: str | None = None):
    """Instantiate the requested LLM provider."""
    if provider_name == "anthropic":
        from .llm.anthropic import AnthropicProvider
        return AnthropicProvider(model=model or "claude-sonnet-4-20250514")
    elif provider_name == "openai":
        from .llm.openai import OpenAIProvider
        return OpenAIProvider(model=model or "gpt-4o")
    elif provider_name == "local":
        from .llm.local import LocalProvider
        return LocalProvider(model=model or "llama3")
    else:
        raise ValueError(f"Unknown provider: {provider_name}")


def run_suggest(
    api_url: str,
    api_key: str,
    project_id: str,
    provider_name: str = "anthropic",
    model: str | None = None,
) -> dict:
    """Execute a single suggest cycle: gather context, call LLM, submit."""
    client = NbenchClient(api_url=api_url, api_key=api_key)
    llm = _get_provider(provider_name, model)

    # Gather project state
    project = client.get_project(project_id)
    experiments = client.get_experiments(project_id, limit=20)
    categories = client.get_categories(project_id)

    try:
        code = client.get_code(project_id)
    except Exception:
        code = ""

    # Build prompt and generate suggestion
    prompt = _build_prompt(project, experiments, categories, code)
    suggestion_text = llm.generate(prompt)

    # Submit suggestion
    result = client.submit_suggestion(project_id, suggestion_text)
    return result
