"""align agent LLM foreground waits with provider budgets

Revision ID: 20260517_agent_runtime_llm_wait_alignment
Revises: 20260517_agent_dynamic_research_prompt_contract
Create Date: 2026-05-17
"""

from __future__ import annotations

import json
import os
from typing import Any

from alembic import op
import sqlalchemy as sa


revision = "20260517_agent_runtime_llm_wait_alignment"
down_revision = "20260517_agent_dynamic_research_prompt_contract"
branch_labels = None
depends_on = None


PROFILE_UPDATES: dict[str, dict[str, Any]] = {
    "pmflow_ai_agent_goal_frame": {
        "provider_timeout_seconds": 45,
        "foreground_wait_seconds": 45,
        "job_timeout_seconds": 55,
        "max_tokens": 1800,
    },
    "pmflow_ai_agent_step_planner": {
        "provider_timeout_seconds": 45,
        "foreground_wait_seconds": 45,
        "job_timeout_seconds": 55,
        "max_tokens": 1800,
    },
    "pmflow_ai_agent_evidence_normalizer": {
        "provider_timeout_seconds": 75,
        "foreground_wait_seconds": 75,
        "job_timeout_seconds": 90,
        "max_tokens": 3200,
    },
    "pmflow_ai_agent_state_critic": {
        "provider_timeout_seconds": 30,
        "foreground_wait_seconds": 30,
        "job_timeout_seconds": 40,
        "max_tokens": 1800,
    },
    "pmflow_ai_agent_answer_composer": {
        "provider_timeout_seconds": 60,
        "foreground_wait_seconds": 60,
        "job_timeout_seconds": 75,
        "max_tokens": 3600,
    },
    "pmflow_ai_agent_answer_reviser": {
        "provider_timeout_seconds": 45,
        "foreground_wait_seconds": 45,
        "job_timeout_seconds": 55,
        "max_tokens": 3000,
    },
}


def _table_exists(bind: Any, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def upgrade() -> None:
    bind = op.get_bind()
    if not _table_exists(bind, "integrations"):
        return
    tenant_id = int(os.getenv("PMFLOW_PLATFORM_RUNTIME_TENANT_ID", "1") or "1")
    row = bind.execute(
        sa.text(
            """
            SELECT id, config FROM integrations
            WHERE tenant_id = :tenant_id
              AND type = 'ai_function_runtime_profiles'
              AND is_active = true
            ORDER BY id DESC
            LIMIT 1
            """
        ),
        {"tenant_id": tenant_id},
    ).mappings().first()
    if row is None:
        return
    config = row.get("config") if isinstance(row.get("config"), dict) else {}
    changed = False
    for stage in ("active", "draft"):
        stage_packet = config.get(stage) if isinstance(config.get(stage), dict) else {}
        profiles = stage_packet.get("profiles") if isinstance(stage_packet.get("profiles"), dict) else {}
        for profile_name, updates in PROFILE_UPDATES.items():
            profile = profiles.get(profile_name)
            if not isinstance(profile, dict):
                continue
            for key, value in updates.items():
                if profile.get(key) != value:
                    profile[key] = value
                    changed = True
    if not changed:
        return
    config["active_version"] = "profile-v20260517-agent-runtime-llm-wait-alignment"
    config["version"] = config["active_version"]
    bind.execute(
        sa.text("UPDATE integrations SET config = CAST(:config AS jsonb) WHERE id = :id"),
        {"id": int(row["id"]), "config": json.dumps(config, ensure_ascii=False)},
    )


def downgrade() -> None:
    return None