"""Cost plan for the residual AI work items, with budget guard (section 18).

Runs BEFORE any real execution: UNIQUE_AI_WORK_ITEMS, token estimates per task type,
estimated cost, and the authorize/refuse verdict against AI_MAX_RUN_COST_USD.
"""
import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from sqlalchemy import select
from central.app.ai.provider import BudgetGuard
from central.app.config import settings
from central.app.db import SessionLocal
from central.app.models import AiWorkItem

TOKENS_PER_ITEM = {  # conservative prompt+schema overhead per single work item
    "INVALID_NCM": (700, 120),
    "CEST": (500, 80),
    "EX_TIPI": (500, 80),
    "PRODUCT_ROLE": (400, 60),
}


def main(output="docs/AI_COST_PLAN_V1.json") -> dict:
    session = SessionLocal()
    items = session.scalars(select(AiWorkItem)).all()
    guard = BudgetGuard(settings.ai_max_run_cost_usd, settings.ai_price_in_per_m, settings.ai_price_out_per_m)
    per_task = {}
    in_total = out_total = 0
    resolved = {"KB_HIT", "KB_EXACT_HIT", "KB_FINGERPRINT_HIT", "HIGH_CONFIDENCE_RETRIEVAL", "SINGLE_OFFICIAL_CANDIDATE", "DONE", "SKIPPED", "CONFLICT"}
    for item in items:
        if item.status in resolved:
            continue  # cache/official candidate resolved — no AI call needed
        tokens_in, tokens_out = TOKENS_PER_ITEM.get(item.task_type, (500, 80))
        bucket = per_task.setdefault(item.task_type, {"items": 0, "input_tokens": 0, "output_tokens": 0, "cells_covered": 0})
        bucket["items"] += 1; bucket["input_tokens"] += tokens_in; bucket["output_tokens"] += tokens_out
        bucket["cells_covered"] += item.cells_covered
        in_total += tokens_in; out_total += tokens_out
    plan = guard.estimate(in_total, out_total)
    report = {
        "artifact": "AI_COST_PLAN_V1",
        "unique_ai_work_items": sum(bucket["items"] for bucket in per_task.values()),
        "per_task_type": per_task,
        "estimated_input_tokens": plan["input_tokens"],
        "estimated_output_tokens": plan["output_tokens"],
        "estimated_cost_usd": plan["estimated_cost"],
        "ai_max_run_cost_usd": guard.max_run_cost_usd,
        "budget_authorized": guard.authorize(plan),
        "provider_configured": settings.ai_provider,
        "model_configured": settings.ai_model,
        "ai_enabled": settings.ai_enabled,
        "max_batch_size": settings.ai_max_batch_size,
        "min_confidence": settings.ai_min_confidence,
    }
    Path(output).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    session.close()
    return report


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", default="/tmp/ai_cost_plan.json")
    args = parser.parse_args()
    print(json.dumps(main(args.output), indent=2))