# -*- coding: utf-8 -*-
"""Week4_Exercise.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1LBlv2wMV5aSZA4Mfnp338rnOdlBzZKHK
"""

import langchain_experimental.tools as tools
print(dir(tools))

# importations
from typing import Annotated, List, Any, Optional, Dict, TypedDict
from typing_extensions import TypedDict
from pydantic import BaseModel, Field, ValidationError
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph
from pydantic import ValidationError
from IPython.display import display
from langgraph.prebuilt import ToolNode
from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRun, DuckDuckGoSearchResults
from langchain_experimental.tools import PythonREPLTool
from langchain.agents import Tool
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.utilities import WikipediaAPIWrapper
import time
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet


import os
import json
import re
import gradio as gr
import requests
import os
import asyncio
import re
import uuid
import tempfile

# load .env
load_dotenv(override=True)

MODEL="claude-haiku-4-5"
BASE_URL="https://api.anthropic.com/v1"

"""Planner Node"""

# Defining the planner schema
class Plan(BaseModel):
    goal: str = Field(..., description="The overall objective")
    steps: List[str] = Field(..., description="Key steps to achieve the goal")

# Planner Prompt
planner_prompt = ChatPromptTemplate.from_messages([
    ("system", """You are a strategic planning agent. Make sure you include all steps needed to achieve the goal.

Return your response as *pure JSON only*.
NO code fences, NO markdown, NO commentary.

Return exactly this format:

{{
  "goal": "<high-level objective>",
  "steps": [
    "<step 1>",
    "<step 2>",
    "<step 3>"
  ]
}}
"""),
    ("human", "Topic: {topic}")
])

# OpenAI-compatible endpoint to Define the planner LLM
planner_llm = ChatOpenAI(
    model=MODEL,
    base_url=BASE_URL,
    api_key=os.getenv("ANTHROPIC_API_KEY"),
    response_format={"type": "json_object"}  # ensures clean JSON
)

#the node Functions
def planner_node(state):
    query = state.get("input", "")
    print("[Planner] Received query:", query)

    prompt = planner_prompt.invoke({"topic": query})
    msg = planner_llm.invoke(prompt)

    # do away with code Fences
    content = msg.content
    content = re.sub(r"```(?:json)?\n", "", content)
    content = re.sub(r"\n```", "", content)
    content = content.strip()

    # Validation
    plan = Plan.model_validate_json(content)
    print("[Planner] Generated plan:", plan)
    return {"plan": plan.dict()}

"""Build the Worker Node"""

# a Wrapper
wiki_wrapper = WikipediaAPIWrapper()

# Tools
search_tool = DuckDuckGoSearchResults()
wiki_tool = WikipediaQueryRun(api_wrapper=wiki_wrapper)
python_tool = PythonREPLTool()

# some Routing
tools = [search_tool, wiki_tool, python_tool]

"""Worker llm"""

worker_llm = ChatOpenAI(
    model=MODEL,
    base_url=BASE_URL,
    api_key=os.getenv("ANTHROPIC_API_KEY")
)

# routing
worker_node = ToolNode(tools=tools)

class SidekickState(TypedDict):
    input: str
    plan: Optional[Dict[str, Any]]
    worker_output: Optional[str]
    review: Optional[Dict[str, Any]]

#worker node
def worker_logic(state: Dict[str, Any]):
    """
    Worker node logic — executes the plan steps using the worker_llm.
    Returns a dictionary with 'worker_output' to feed back into the graph.
    """

    plan = state.get("plan")
    if not plan:
        return {"worker_output": "No plan available."}

    # planning structue
    if isinstance(plan, dict):
        goal = plan.get("goal", "")
        steps = plan.get("steps", [])
    elif isinstance(plan, list):
        goal = "Execute each step in the provided plan."
        steps = plan
    else:
        goal = getattr(plan, "goal", "")
        steps = getattr(plan, "steps", [])

    print(f"[Worker] Executing plan for goal: {goal}")
    summary = ""

    for i, step in enumerate(steps, start=1):
        print(f"[Worker] Step {i}: {step}")

        prompt = (
            f"You are a business research analyst.\n"
            f"Step {i}: {step}\n"
            f"Goal: {goal}\n"
            f"Please produce a clear, concise actionable output."
        )

        result, content = None, None
        for attempt in range(2):
            try:
                result = worker_llm.invoke(prompt)
                content = getattr(result, "content", None)

                # Handle nested Anthropic / OpenAI structure
                if not content:
                    try:
                        content = result.generations[0][0].message.content
                    except Exception:
                        content = str(result)

                if content:
                    break
            except Exception as e:
                print(f"[Worker] Retry {attempt+1} failed: {e}")
                time.sleep(2)

        summary += f"\nStep {i}: {step}\nResult: {content or 'No response received.'}\n"

    print("[Worker] Completed all steps.")
    return {"worker_output": summary}

"""Reviewer Node"""

# Review schema
class Review(BaseModel):
    clarity_score: int = Field(..., description="Score 1–10 for how clear the report is")
    completeness_score: int = Field(..., description="Score 1–10 for how complete the analysis is")
    feedback: str = Field(..., description="Suggestions for improving the report")
    improved_summary: str = Field(..., description="An improved or polished version of the report")

# llm
reviewer_llm = ChatOpenAI(
    model=MODEL,
    base_url=BASE_URL,
    api_key=os.getenv("ANTHROPIC_API_KEY")
)

#review node
def review_logic(state: SidekickState):
    output = state.get("worker_output", "")
    if not output:
        return {"review": {"feedback": "No worker output to review."}}

    print("[Reviewer] Evaluating output...")
    feedback = reviewer_llm.invoke(
        f"""Review the following consultant report output:

{output}

Provide your feedback in this structured format:
Clarity Score (1–10):
Completeness Score (1–10):
Suggestions for improvement:
Improved Summary:"""
    )

    content = getattr(feedback, "content", None)
    if not content and hasattr(feedback, "generations"):
        content = feedback.generations[0][0].message.content

    return {"review": {"feedback": content or "No feedback generated."}}

"""Build and Connect the LangGraph"""

class SidekickState(TypedDict):
    input: str
    plan: Optional[Dict[str, Any]]
    worker_output: Optional[str]
    #review: Optional[Dict[str, Any]]
    review: Optional[str]

# the graph
graph = StateGraph(SidekickState)

planner_logic = planner_node

# directly register nodes
graph.add_node("planner", planner_logic)
graph.add_node("worker", worker_logic)
graph.add_node("reviewer", review_logic)

# Flow
graph.add_edge(START, "planner")
graph.add_edge("planner", "worker")
graph.add_edge("worker", "reviewer")
graph.add_edge("reviewer", END)

# Chekpointing
memory = MemorySaver()

consultant_graph = graph.compile(checkpointer=memory)

print(" ConsultantSidekick Graph compiled successfully!")

"""Pushover"""

pushover_token = os.getenv("PUSHOVER_TOKEN")
pushover_user = os.getenv("PUSHOVER_USER")
pushover_url = "https://api.pushover.net/1/messages.json"

def push(text: str):
    if pushover_token and pushover_user:
        requests.post(
            pushover_url,
            data={"token": pushover_token, "user": pushover_user, "message": text},
        )
    else:
        print(" No Pushover credentials found. Skipping notification.")

"""Gradio"""

async def run_consultant_sidekick(user_input, history):
    init_state = SidekickState(input=user_input)

    # Checkpointing
    thread_id = f"session_{uuid.uuid4()}"
    checkpoint_ns = "consultant_sidekick"
    checkpoint_id = str(uuid.uuid4())

    # config nesting
    result = await consultant_graph.ainvoke(
        input=init_state,
        config={
            "configurable": {
                "thread_id": thread_id,
                "checkpoint_ns": checkpoint_ns,
                "checkpoint_id": checkpoint_id,
            }
        }
    )

    # plan = result.plan.dict() if getattr(result, "plan", None) else {}
    # review = result.review.dict() if getattr(result, "review", None) else {}
    plan = result.get("plan", {})
    review = result.get("review", {})


    output_text = (
        f"**Plan:** {plan.get('steps', [])}\n\n"
        #f"**Output:** {getattr(result, 'worker_output', 'No output')}\n\n"
        f"**Output:** {result.get('worker_output', 'No output')}\n\n"
        f"**Reviewer Feedback:** {review.get('feedback', 'No feedback')}"
    )

    # push notification
    push("Hey there, check your report from sidekick!")

    # history
    history.append((user_input, output_text))
    return history, ""


"""pdfs"""

def clean_text(text: str) -> str:
    # bold/italic
    text = re.sub(r"(\*\*|__)(.*?)\1", r"\2", text)
    text = re.sub(r"(\*|_)(.*?)\1", r"\2", text)
    # headings
    text = re.sub(r"#+\s*", "", text)
    # tale handlig
    text = re.sub(r"^\|.*\|$", "", text, flags=re.MULTILINE)
    # newlines
    text = text.replace("<br/>", "\n")
    return text.strip()

def create_pdf(title, content):
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=A4)
    styles = getSampleStyleSheet()
    story = []

    story.append(Paragraph(f"<b>{title}</b>", styles["Title"]))
    story.append(Spacer(1, 12))

    if isinstance(content, list):
        for i, point in enumerate(content, 1):
            clean_point = clean_text(str(point))
            story.append(Paragraph(f"{i}. {clean_point}", styles["BodyText"]))
            story.append(Spacer(1, 6))
    else:
        clean_content = clean_text(str(content))
        story.append(Paragraph(clean_content.replace("\n", "<br/>"), styles["BodyText"]))

    doc.build(story)
    buffer.seek(0)
    return buffer

def download_plan(plan):
    pdf_buffer = create_pdf("Consultant Plan", plan["steps"])
    return pdf_buffer, "plan_report.pdf"

def download_output(output):
    pdf_buffer = create_pdf("Detailed Execution Report", output)
    return pdf_buffer, "output_report.pdf"

def download_review(review):
    pdf_buffer = create_pdf("Evaluator Feedback", review)
    return pdf_buffer, "review_report.pdf"

#gradio

def safe_update(value=None, visible=None):
    """Works across all Gradio 4.x versions."""
    if hasattr(gr, "update"):
        return gr.update(value=value, visible=visible)
    return {"value": value, "visible": visible}

with gr.Blocks() as app:
    gr.Markdown("## Consultant-Ask anything and get a report")

    with gr.Row():
        query = gr.Textbox(
            label="what would you like to consult about?",
            placeholder="e.g., Goat Farming in Eastern, kenya",
            lines=1
        )
        run_btn = gr.Button("Consult?")

    status = gr.Markdown("", visible=False)

    # hid buttons at firs
    with gr.Row():
        download_plan_btn = gr.File(label="See the report Plan", visible=False)
        download_output_btn = gr.File(label="See the actual Report", visible=False)
        download_review_btn = gr.File(label="See why the report is good", visible=False)

    async def run_all(query):
        yield (
            safe_update(value="**Working on the report, please bear with me**", visible=True),
            safe_update(), safe_update(), safe_update()
        )

        thread_id = f"session_{uuid.uuid4()}"
        checkpoint_ns = "consultant_sidekick"
        checkpoint_id = str(uuid.uuid4())

        result = await consultant_graph.ainvoke(
            input={"input": query},
            config={
                "configurable": {
                    "thread_id": thread_id,
                    "checkpoint_ns": checkpoint_ns,
                    "checkpoint_id": checkpoint_id
                }
            }
        )

        plan = result.get("plan", {})
        worker = result.get("worker_output", "")
        review = result.get("review", {})

        # pdf saving
        def save_temp_pdf(buffer):
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
            tmp.write(buffer.read())
            tmp.close()
            buffer.seek(0)
            return tmp.name

        plan_pdf_path = save_temp_pdf(download_plan(plan)[0])
        output_pdf_path = save_temp_pdf(download_output(worker)[0])
        review_pdf_path = save_temp_pdf(download_review(review)[0])

        # --- Yield completed state ---
        yield (
            safe_update(value="**Done! get your report.**", visible=True),
            safe_update(value=plan_pdf_path, visible=True),
            safe_update(value=output_pdf_path, visible=True),
            safe_update(value=review_pdf_path, visible=True),
        )

    run_btn.click(
        run_all,
        inputs=query,
        outputs=[
            status,
            download_plan_btn,
            download_output_btn,
            download_review_btn,
        ],
        show_progress=False,
    )

# --- Launch ---
app.launch(inbrowser=True)