"""
AI-powered lineage explanation for a revenue pipeline.

Pipeline: prices * quantities + bonus = total_revenue

Runs the pipeline under preservation mode "FULL" so all intermediate tables
are preserved for debugging. The debug agent uses its tool loop to inspect
tables and trace the computation graph.
"""

import asyncio
import logging
import os

from aaiclick.ai.agents.debug_agent import debug_result
from aaiclick.ai.agents.lineage_agent import explain_lineage
from aaiclick.data.data_context import create_object_from_value
from aaiclick.data.object import Object
from aaiclick.oplog.lineage import lineage_context, oplog_subgraph
from aaiclick.orchestration import (
    JOB_COMPLETED,
    PRESERVATION_FULL,
    JobStatus,
    PreservationMode,
    ajob_test,
    get_tasks_for_job,
    job,
    task,
    tasks_list,
)
from aaiclick.orchestration.orch_context import orch_context

from .report import print_report

logger = logging.getLogger(__name__)


@task
async def create_prices() -> Object:
    return await create_object_from_value(
        [10.0, 20.0, 30.0, 40.0, 50.0],
        name="basic_lineage_prices",
        aai_id=True,
    )


@task
async def create_quantities() -> Object:
    return await create_object_from_value(
        [2.0, 3.0, 1.0, 5.0, 4.0],
        name="basic_lineage_quantities",
        aai_id=True,
    )


@task
async def create_bonus() -> Object:
    return await create_object_from_value(
        [5.0, 5.0, 5.0, 5.0, 5.0],
        name="basic_lineage_bonus",
        aai_id=True,
    )


@task
async def compute_revenue(prices: Object, quantities: Object) -> Object:
    return (prices * quantities).as_("revenue")


@task
async def add_bonus(revenue: Object, bonus: Object) -> Object:
    return (revenue + bonus).as_("total")


@job("revenue_pipeline")
def revenue_pipeline():
    prices = create_prices()
    quantities = create_quantities()
    bonus = create_bonus()
    revenue = compute_revenue(prices=prices, quantities=quantities)
    total = add_bonus(revenue=revenue, bonus=bonus)
    return tasks_list(prices, quantities, bonus, revenue, total)


async def main():
    async with orch_context():
        pipeline = await revenue_pipeline(
            preservation_mode=PRESERVATION_FULL,
        )
        await ajob_test(pipeline)
        assert pipeline.status == JOB_COMPLETED, f"Job failed: {pipeline.error}"

        tasks = await get_tasks_for_job(pipeline.id)
        target_table = next(t for t in tasks if t.name == "add_bonus").result["table"]
        source_table = next(t for t in tasks if t.name == "create_prices").result["table"]

        explanation: str | None = None
        debug_answer: str | None = None
        async with lineage_context():
            backward_graph, forward_graph = await asyncio.gather(
                oplog_subgraph(target_table, direction="backward"),
                oplog_subgraph(source_table, direction="forward"),
            )
            if os.environ.get("AAICLICK_AI_API_KEY"):
                explanation = await explain_lineage(
                    target_table,
                    question="How was this table produced? What arithmetic was applied?",
                    graph=backward_graph,
                )
                debug_answer = await debug_result(
                    target_table,
                    question=(
                        "Which output row has the highest value, and which input "
                        "rows drove it? Use the tools to inspect the tables."
                    ),
                    graph=backward_graph,
                )
            else:
                logger.warning("AAICLICK_AI_API_KEY is not set; skipping AI lineage explanation and debug agent.")

        print_report(
            tasks=tasks,
            target_table=target_table,
            backward_graph=backward_graph,
            forward_graph=forward_graph,
            source_table=source_table,
            explanation=explanation,
            debug_answer=debug_answer,
        )
