"""
09_self_healing_pipeline.py — Self-healing pipeline with retry logic and validator loop

Tests the Failure History + retry patterns. A generator agent produces output,
a validator agent checks it, and if validation fails the generator retries
(up to 3 times). The graph shows multiple generator instances if retries occur.

Pipeline:
  task_planner → generator (attempt 1) → validator
                                            ↓ (fail)
                           generator (attempt 2) → validator
                                            ↓ (fail / pass)
                           generator (attempt 3) → validator
                                            ↓ (pass)
                           → quality_gate → finalizer

WHAT THE DASHBOARD SHOULD SHOW:
  Live Graph:
    • task_planner → generator_attempt_1 → validator_1
    • If retry: validator_1 → generator_attempt_2 → validator_2 → ...
    • quality_gate → finalizer at the end
    • Multiple "generator" nodes visible (one per attempt) — same function, unique agent_ids

  State View:
    • Each validator step shows the validation result and what FAILED
    • Generator attempts show how the prompt changes each retry
    • quality_gate INPUT shows final passed output

  Time Travel:
    • Rewind to see earlier (failed) generator attempts
    • Compare attempt_1 output vs attempt_2 output to see improvement

  Token Dashboard:
    • Total tokens increase with each retry — visible in line chart
    • quality_gate bar is small (just a check, no LLM)

  Run Comparison:
    • Shows all generator attempts as separate rows
    • Failed validators show as FAILED rows (red) in the table

  Failure History:
    • If all 3 attempts fail, the run is auto-saved to Failure History
    • If any attempt passes, the failure history entry is NOT created
"""
import asyncio
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from agentradar import radar, reset_run

load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    temperature=0.5,
)

MAX_RETRIES = 3


@radar(tag="task_planner", llm_input="Break the task into clear requirements and success criteria")
async def task_planner(task: str) -> dict:
    """Step 1: Define the task requirements and what 'done' looks like."""
    message = llm.invoke(
        f"Break this task into 3 clear requirements and define success criteria (measurable):\n{task}"
    )
    return {
        "task": task,
        "requirements": message.content,
        "max_retries": MAX_RETRIES,
    }


@radar(tag="generator", llm_input="Generate output satisfying all requirements")
async def generator(task_plan: dict, attempt: int, feedback: str = "") -> dict:
    """
    Generate a solution. Called multiple times if validation fails.
    Each call creates a NEW node in the graph (unique agent_id per call).
    """
    feedback_section = f"\n\nPrevious attempt failed. Feedback: {feedback}" if feedback else ""
    message = llm.invoke(
        f"Task requirements:\n{task_plan['requirements']}\n"
        f"Attempt {attempt}/{MAX_RETRIES}.{feedback_section}\n\n"
        f"Generate a complete solution that satisfies ALL requirements. "
        f"Format as a numbered list of deliverables."
    )
    return {
        "attempt": attempt,
        "output": message.content,
        "task": task_plan["task"],
    }


@radar(tag="validator", llm_input=["generator"])
async def validator(generated: dict, requirements: str) -> dict:
    """
    Validate that the generated output meets all requirements.
    Returns passed=True or passed=False with detailed feedback.
    Each call is a separate node in the graph.
    """
    message = llm.invoke(
        f"Requirements:\n{requirements}\n\n"
        f"Generated output (attempt {generated['attempt']}):\n{generated['output']}\n\n"
        f"Does this output satisfy ALL requirements? Answer:\n"
        f"PASS or FAIL\n"
        f"Missing items (if FAIL): list them\n"
        f"Confidence (0-100%): give a number"
    )
    response = message.content
    passed = "PASS" in response.upper() and "FAIL" not in response.upper()[:20]
    return {
        "attempt": generated["attempt"],
        "passed": passed,
        "feedback": response,
        "output": generated["output"],
    }


@radar(tag="quality_gate", llm_input=["validator"])
async def quality_gate(validation_result: dict) -> dict:
    """
    Final gate: only passes if validation succeeded.
    No LLM call — pure logic gate (tests 0-token agents in dashboard).
    """
    if not validation_result["passed"]:
        raise RuntimeError(
            f"Quality gate FAILED after {validation_result['attempt']} attempts. "
            f"Feedback: {validation_result['feedback'][:200]}"
        )
    return {
        "approved": True,
        "attempts_needed": validation_result["attempt"],
        "final_output": validation_result["output"],
    }


@radar(tag="finalizer", llm_input=["quality_gate"])
async def finalizer(approved: dict) -> dict:
    """Final step: format and deliver the approved output."""
    message = llm.invoke(
        f"Format this approved solution as a clean, professional deliverable:\n{approved['final_output']}"
    )
    return {
        "status": "delivered",
        "attempts_taken": approved["attempts_needed"],
        "deliverable": message.content,
    }


async def main():
    reset_run()
    task = "Write a Python function to securely hash passwords using bcrypt with salt"

    print("=== Self-Healing Pipeline (Retry Loop) ===\n")

    print("[1] Planning task...")
    plan = await task_planner(task)
    print(f"  → Requirements defined\n")

    validation_result = None
    feedback = ""

    for attempt in range(1, MAX_RETRIES + 1):
        print(f"[Attempt {attempt}/{MAX_RETRIES}] Generating...")
        generated = await generator(plan, attempt, feedback)
        print(f"  → Generated {len(generated['output'].split())} words\n")

        print(f"[Attempt {attempt}/{MAX_RETRIES}] Validating...")
        validation_result = await validator(generated, plan["requirements"])
        status = "✓ PASS" if validation_result["passed"] else "✗ FAIL"
        print(f"  → {status}\n")

        if validation_result["passed"]:
            break
        else:
            feedback = validation_result["feedback"]
            if attempt < MAX_RETRIES:
                print(f"  [Retry] Incorporating feedback for attempt {attempt + 1}...\n")

    print("[Quality Gate] Checking final output...")
    try:
        approved = await quality_gate(validation_result)
        print(f"  → Approved after {approved['attempts_needed']} attempt(s)\n")

        print("[Finalizer] Delivering...")
        final = await finalizer(approved)
        print(f"  → Delivered! Attempts taken: {final['attempts_taken']}")
        print(f"  → Preview: {final['deliverable'][:150]}...")
    except RuntimeError as e:
        print(f"  ✗ {e}")
        print("  → Run saved to Failure History")

    print(f"\n✓ Dashboard: {attempt} generator node(s), retry loop visible, quality_gate = 0 tokens")


if __name__ == "__main__":
    asyncio.run(main())
    input("\nPress Enter to view dashboard (http://localhost:7111) or Ctrl+C to exit...")
