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

Automatically generated by Colab.

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

Revised notebook. Now using an updated model: claude-haiku-4-5 (and cheaper: as was recommended).
The notebook is also updated to use the OpenAI format as in Guide 9 of Ed's Guide (many thanks Ed for keeping me on track on this)
"""

# all rqments needed in the exercise.
!pip install openai>=1.40.0 gradio reportlab python-dotenv -q

# all imports needed throught the exercise
import os, json, asyncio, tempfile
import nest_asyncio
from datetime import datetime
from dotenv import load_dotenv
from pydantic import BaseModel
from typing import Any, Dict, List
from openai import OpenAI
import gradio as gr
from reportlab.lib.pagesizes import A4
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet
import re

from google.colab import userdata
ANTHROPIC_API_KEY = userdata.get("ANTHROPIC_API_KEYS")

# Anthropic’s OpenAI-compatible endpoint
ANTHROPIC_BASE_URL = "https://api.anthropic.com/v1"

# to confirm the Loading was successful
if not ANTHROPIC_API_KEY:
    raise ValueError("Missing Anthropic API key..")

#instance to initalize the openAI-compatible client
client = OpenAI(
    base_url=ANTHROPIC_BASE_URL,
    api_key=ANTHROPIC_API_KEY
)

#recommended anthropic model (cheaper +updated)
MODEL = "claude-haiku-4-5"

# confirm the chat completions work on this
completion = client.chat.completions.create(
    model=MODEL,
    messages=[
        {"role": "system", "content": "You are Claude, a helpful assistant."},
        {"role": "user", "content": "Give me a 1-line summary of AI trends in 2025."}
    ]
)
response=completion.choices[0].message.content
print(response)

#call_anthropic helper
def call_anthropic(prompt: str, temperature: float = 0.3, max_tokens: int = 1000):
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {"role": "system", "content": "You are a helpful research assistant that always responds in valid JSON."},
            {"role": "user", "content": prompt}
        ],
        temperature=temperature,
        max_tokens=max_tokens
    )
    return completion.choices[0].message.content.strip()

#our Clarifier agent
class ClarifierOutput(BaseModel):
    questions: List[str]
    refined_query: str

class ClarifierAgent:
    async def run(self, user_query: str) -> ClarifierOutput:
        prompt = f"""
You are a Clarifier Agent. Given a user's research query, return *only* valid JSON:

{{
  "questions": ["question1", "question2", "question3"],
  "refined_query": "refined query"
}}

User query:
\"\"\"{user_query}\"\"\"
"""
        raw = await asyncio.to_thread(call_anthropic, prompt)
        try:
            data = json.loads(raw)
        except Exception:
            start, end = raw.find("{"), raw.rfind("}") + 1
            data = json.loads(raw[start:end])
        return ClarifierOutput(**data)

class KeyFinding(BaseModel):
    title: str
    detail: str

class ResearcherOutput(BaseModel):
    short_summary: str
    key_findings: List[KeyFinding]
    suggested_sources: List[str]

#now the Researcher agent
class ResearcherAgent:
    async def run(self, clarified_query: str) -> ResearcherOutput:
        prompt = f"""
You are a Research Agent. Conduct high-level research on this query and return JSON only.

Clarified query:
\"\"\"{clarified_query}\"\"\"

JSON format:
{{
  "short_summary": "...",
  "key_findings": [
    {{"title": "Finding 1", "detail": "Details"}},
    {{"title": "Finding 2", "detail": "Details"}}
  ],
  "suggested_sources": ["Source 1", "Source 2"]
}}
"""
        raw = await asyncio.to_thread(call_anthropic, prompt)
        try:
            data = json.loads(raw)
        except Exception:
            start, end = raw.find("{"), raw.rfind("}") + 1
            data = json.loads(raw[start:end])

        key_findings = [KeyFinding(**f) for f in data.get("key_findings", [])]
        return ResearcherOutput(
            short_summary=data.get("short_summary", ""),
            key_findings=key_findings,
            suggested_sources=data.get("suggested_sources", []),
        )

# evaluation agent
class EvaluationOutput(BaseModel):
    clarity: str
    accuracy: str
    completeness: str
    strengths: List[str]
    weaknesses: List[str]
    recommendations: List[str]

class EvaluatorAgent:
    async def run(self, summary: str, findings: List[KeyFinding], refined_query: str) -> EvaluationOutput:
        findings_text = "\n".join([f"- {f.title}: {f.detail}" for f in findings])
        prompt = f"""
You are an expert Evaluator. Assess the quality of this research summary.

Return JSON only in this format:
{{
  "clarity": "...",
  "accuracy": "...",
  "completeness": "...",
  "strengths": ["..."],
  "weaknesses": ["..."],
  "recommendations": ["..."]
}}

Findings:
{findings_text}

Summary:
{summary}

Refined Query:
{refined_query}
"""
        raw = await asyncio.to_thread(call_anthropic, prompt)
        try:
            data = json.loads(raw)
        except Exception:
            start, end = raw.find("{"), raw.rfind("}") + 1
            data = json.loads(raw[start:end])
        return EvaluationOutput(**data)

# researchmanager
class ResearchManager:
    def __init__(self):
        self.clarifier = ClarifierAgent()
        self.researcher = ResearcherAgent()
        self.evaluator = EvaluatorAgent()

    async def run(self, query: str) -> dict:
        print("\n Clarifying query...")
        clar = await self.clarifier.run(query)
        print("Clarified:", clar.refined_query)

        print("\n Conducting research...")
        research = await self.researcher.run(clar.refined_query)

        print("\n Evaluating results...")
        eval = await self.evaluator.run(research.short_summary, research.key_findings, clar.refined_query)

        return {
            "original_query": query,
            "clarifying_questions": clar.questions,
            "refined_query": clar.refined_query,
            "research_summary": research.short_summary,
            "key_findings": [f.dict() for f in research.key_findings],
            "suggested_sources": research.suggested_sources,
            "evaluation": eval.dict(),
        }

#Pdf generation
def generate_pdf(report, prefix="report"):
    safe_prefix = "_".join(prefix.lower().split())[:50]
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"deep_research_{safe_prefix}_{timestamp}.pdf"
    path = os.path.join(tempfile.gettempdir(), filename)
    doc = SimpleDocTemplate(path, pagesize=A4)
    styles = getSampleStyleSheet()
    elems = []

    def section(title, content):
        elems.append(Paragraph(f"<b>{title}</b>", styles["Heading3"]))
        elems.append(Spacer(1, 8))
        elems.append(Paragraph(content.replace("\n", "<br/>"), styles["Normal"]))
        elems.append(Spacer(1, 12))

    elems.append(Paragraph("<b>AI Deep Research Report</b>", styles["Title"]))
    elems.append(Spacer(1, 12))
    section("Original Query", report["original_query"])
    section("Refined Query", report["refined_query"])
    section("Clarifying Questions", "\n".join(report["clarifying_questions"]))
    section("Summary", report["research_summary"])

    findings = "\n".join([f"- {f['title']}: {f['detail']}" for f in report["key_findings"]])
    section("Key Findings", findings)

    eval = report["evaluation"]
    eval_text = f"""
Clarity: {eval['clarity']}
Accuracy: {eval['accuracy']}
Completeness: {eval['completeness']}

Strengths:
{chr(10).join(eval['strengths'])}

Weaknesses:
{chr(10).join(eval['weaknesses'])}

Recommendations:
{chr(10).join(eval['recommendations'])}
"""
    section("Evaluation", eval_text)
    doc.build(elems)
    return path

#workflow
async def process_query(query: str):
    manager = ResearchManager()
    result = await manager.run(query)
    pdf_path = generate_pdf(result, prefix=query)
    json_path = os.path.join(tempfile.gettempdir(), f"{query.replace(' ', '_')}_data.json")
    with open(json_path, "w") as f:
        json.dump(result, f, indent=2)
    return pdf_path, json_path, " Report generated successfully!"

#The UX
nest_asyncio.apply()
async def on_generate_click(q):
    # regex to have a good file name
    safe_prefix = re.sub(r'[^a-zA-Z0-9_-]', '_', q)[:50]
    pdf_path, json_path, msg = await process_query(q)

    if pdf_path and json_path:
        # up date the output if successful
        return (
            pdf_path,
            json_path,
            f"<p style='color:green;'>{msg}</p>",
            gr.update(interactive=True),
            gr.update(interactive=True)
        )
    else:
        #if not successful, how to update
        return (
            None,
            None,
            f"<p style='color:red;'>{msg}</p>",
            gr.update(interactive=False),
            gr.update(interactive=False)
        )

#The UX now with GRadio
with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    # app header
    with gr.Row():
        gr.Markdown(
            """
            <h1 style='text-align:center; font-weight:700;'>
                AI Research Report Generator
            </h1>
            <p style='text-align:center; font-size:16px; color:gray;'>
                Enter your research topic or question to automatically generate a structured research report with clarifications, analysis, and evaluation.
            </p>
            """,
            elem_id="header"
        )

    # Input field
    query_input = gr.Textbox(
        label="Research Topic or Question",
        placeholder="e.g., Secrets of becoming an AI expert in 2025",
        lines=3
    )

    # Generate button
    generate_btn = gr.Button("Generate Report", variant="primary")

    # Outputs
    with gr.Row():
        pdf_output = gr.File(label="Download PDF Report", interactive=False)
        json_output = gr.File(label="Download JSON Data", interactive=False)

    # Status message
    status_msg = gr.Markdown("", elem_id="status_text")

    # Connect async callback
    generate_btn.click(
        fn=on_generate_click,
        inputs=query_input,
        outputs=[pdf_output, json_output, status_msg, pdf_output, json_output]
    )

demo.launch()