main.py

"""
Intern Problem-Solving API
Multi-agent FastAPI backend for structured problem analysis and solution generation.
Agents: Analyst → Root Cause → Solution Brainstorm → Action Planner → PDF Generator
"""

import io
import json
import os
import re
from datetime import datetime
from typing import Optional

import anthropic
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import ParagraphStyle, getSampleStyleSheet
from reportlab.lib.units import mm
from reportlab.platypus import (
    HRFlowable,
    Paragraph,
    SimpleDocTemplate,
    Spacer,
    Table,
    TableStyle,
)

# ── App Setup ─────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Intern Problem-Solving API",
    description="Multi-agent system: Analyst → Root Cause → Solutions → Action Plan",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", ""))


# ── PDF Text Extraction ───────────────────────────────────────────────────────

def decode_content(raw: str) -> str:
    """
    If the client sent a PDF as __PDF_BASE64__<data>, decode and extract text.
    Otherwise return the string unchanged.
    """
    PREFIX = "__PDF_BASE64__"
    if not raw.startswith(PREFIX):
        return raw

    import base64
    from pypdf import PdfReader

    b64 = raw[len(PREFIX):]
    try:
        pdf_bytes = base64.b64decode(b64)
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid base64 PDF data.")

    try:
        reader = PdfReader(io.BytesIO(pdf_bytes))
        pages = []
        for page in reader.pages:
            text = page.extract_text()
            if text:
                pages.append(text.strip())
        extracted = "\n\n".join(pages).strip()
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Could not read PDF: {e}")

    if not extracted:
        raise HTTPException(
            status_code=400,
            detail="PDF appears to be scanned/image-based — no text found. Please paste the text manually.",
        )
    return extracted


# ── Request / Response Models ─────────────────────────────────────────────────
class ProblemInput(BaseModel):
    content: str
    intern_name: Optional[str] = "Intern"
    intern_role: Optional[str] = "AI Developer Intern"
    intern_goal: Optional[str] = ""


class AgentOutput(BaseModel):
    agent: str
    output: str


class FullAnalysis(BaseModel):
    problem_statement: str
    root_causes: str
    solutions: str
    action_plan: str
    thinking_feedback: str


# ── Agent Definitions ─────────────────────────────────────────────────────────

AGENT_ANALYST = """You are the Problem Analyst Agent.
Your ONLY job: read the input and produce a crisp, structured problem statement.

Output format (use these exact headers):
## Core Problem
One clear sentence: who has what problem, in what context.

## Key Pain Points
- Bullet each distinct pain point (max 5)

## Stakeholders
- Who is affected and how

## Known vs Unknown
- Known: what facts are clear
- Unknown: what gaps exist

Keep it factual. No solutions yet. Max 250 words."""


AGENT_ROOT_CAUSE = """You are the Root Cause Analysis Agent.
You receive a problem statement. Your job: find WHY it exists.

Output format:
## Root Cause Analysis

### Immediate Cause
What is the visible trigger of the problem?

### Underlying Causes
Break down causes across three lenses:
- **Technical**: systems, tools, architecture issues
- **Process**: workflow, communication, or procedural gaps
- **People/Skills**: knowledge gaps, habits, capacity issues

### The Real Root Cause
One sentence: the deepest cause everything traces back to.

Be specific. Use "because" chains to trace causes deeper. Max 200 words."""


AGENT_SOLUTIONS = """You are the Solution Brainstorm Agent.
You receive a problem + root cause. Generate diverse, creative solutions.

Output format:
## Solution Ideas

### Quick Wins (Do this week)
1. **[Name]** — What it is + why it helps

### Medium-Term Fixes (Do this month)
2. **[Name]** — What it is + why it helps
3. **[Name]** — What it is + why it helps

### Strategic / Long-Term
4. **[Name]** — What it is + why it helps
5. **[Name]** — What it is + why it helps

### Unconventional / Creative
6. **[Name]** — Think outside the box
7. **[Name]** — Wildcard idea

For each idea: name it, describe it in 1-2 sentences, state the trade-off.
Think across: AI tools, process redesign, automation, collaboration, education.
Max 300 words."""


AGENT_ACTION_PLANNER = """You are the Action Planner Agent.
You receive the full analysis. Your job: give the intern 3 concrete next actions.

Output format:
## Your Next Steps

### Action 1: [Do This Today]
**What exactly**: One sentence instruction
**How**: Step-by-step (3-4 steps max)
**Success looks like**: How you'll know it worked
**Time needed**: X hours

### Action 2: [Do This Week]
**What exactly**: One sentence instruction
**How**: Step-by-step (3-4 steps max)
**Success looks like**: How you'll know it worked
**Time needed**: X hours

### Action 3: [Do This Month]
**What exactly**: One sentence instruction
**How**: Step-by-step (3-4 steps max)
**Success looks like**: How you'll know it worked
**Time needed**: X hours

Be specific enough that the intern can start immediately. No vague advice."""


AGENT_THINKING_COACH = """You are the Thinking Coach Agent.
You are an encouraging but honest coach helping an intern grow.
You receive the original problem input + full analysis.

Output format:
## Thinking Feedback

### What You Got Right
- Specific things in how the problem was framed that show good thinking

### Blind Spots to Watch
- Where the framing was shallow or missing something important
- Specific examples only — no generic observations

### Are You Thinking Like a Problem Solver or Task Executor?
One honest assessment with evidence from their input.

### One Big Shift
The single most important mindset or approach shift that will help this intern most.

### For Your Next Meeting
3 specific things to do differently next time you face this type of problem.

Keep it encouraging but honest. Max 250 words."""


# ── Core Agent Runner ─────────────────────────────────────────────────────────

def run_agent(system_prompt: str, user_content: str, max_tokens: int = 800) -> str:
    """Run a single agent and return its text output."""
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=max_tokens,
        system=system_prompt,
        messages=[{"role": "user", "content": user_content}],
    )
    return response.content[0].text


# ── Pipeline ──────────────────────────────────────────────────────────────────

def run_pipeline(content: str, name: str, role: str, goal: str) -> FullAnalysis:
    """Run all 5 agents in sequence, passing outputs forward."""

    context_header = f"""
Intern Name: {name}
Intern Role: {role}
Current Goal: {goal if goal else 'Not specified'}

--- INPUT ---
{content[:8000]}
--- END INPUT ---
"""

    # Agent 1: Analyst
    problem_statement = run_agent(
        AGENT_ANALYST,
        f"Analyze this problem:\n{context_header}",
        max_tokens=600,
    )

    # Agent 2: Root Cause (receives problem statement)
    root_causes = run_agent(
        AGENT_ROOT_CAUSE,
        f"Problem Statement:\n{problem_statement}\n\nOriginal input context:\n{content[:3000]}",
        max_tokens=500,
    )

    # Agent 3: Solutions (receives problem + root causes)
    solutions = run_agent(
        AGENT_SOLUTIONS,
        f"Problem Statement:\n{problem_statement}\n\nRoot Causes:\n{root_causes}",
        max_tokens=700,
    )

    # Agent 4: Action Planner (receives everything so far)
    action_plan = run_agent(
        AGENT_ACTION_PLANNER,
        f"""Intern Role: {role}\nIntern Goal: {goal}

Problem Statement:\n{problem_statement}

Root Causes:\n{root_causes}

Solutions:\n{solutions}""",
        max_tokens=700,
    )

    # Agent 5: Thinking Coach (sees original input + full analysis)
    thinking_feedback = run_agent(
        AGENT_THINKING_COACH,
        f"""Original Input from Intern:\n{content[:3000]}

Problem Analysis:\n{problem_statement}

Root Causes:\n{root_causes}""",
        max_tokens=600,
    )

    return FullAnalysis(
        problem_statement=problem_statement,
        root_causes=root_causes,
        solutions=solutions,
        action_plan=action_plan,
        thinking_feedback=thinking_feedback,
    )


# ── PDF Generator ─────────────────────────────────────────────────────────────

def strip_markdown(text: str) -> str:
    """Remove markdown bold/italic markers for plain PDF text."""
    text = re.sub(r"\*\*(.*?)\*\*", r"\1", text)
    text = re.sub(r"\*(.*?)\*", r"\1", text)
    return text


def build_pdf(analysis: FullAnalysis, name: str, role: str) -> bytes:
    """Generate a professional PDF report from the analysis."""
    buf = io.BytesIO()

    # ── Page setup
    doc = SimpleDocTemplate(
        buf,
        pagesize=A4,
        rightMargin=20 * mm,
        leftMargin=20 * mm,
        topMargin=22 * mm,
        bottomMargin=22 * mm,
    )

    styles = getSampleStyleSheet()
    W = A4[0] - 40 * mm  # usable width

    # ── Custom styles
    s_title = ParagraphStyle(
        "title",
        parent=styles["Normal"],
        fontSize=22,
        fontName="Helvetica-Bold",
        textColor=colors.HexColor("#1a1a2e"),
        spaceAfter=4,
    )
    s_sub = ParagraphStyle(
        "sub",
        parent=styles["Normal"],
        fontSize=11,
        fontName="Helvetica",
        textColor=colors.HexColor("#6b7280"),
        spaceAfter=12,
    )
    s_section = ParagraphStyle(
        "section",
        parent=styles["Normal"],
        fontSize=13,
        fontName="Helvetica-Bold",
        textColor=colors.HexColor("#2563eb"),
        spaceBefore=16,
        spaceAfter=6,
    )
    s_h3 = ParagraphStyle(
        "h3",
        parent=styles["Normal"],
        fontSize=11,
        fontName="Helvetica-Bold",
        textColor=colors.HexColor("#374151"),
        spaceBefore=8,
        spaceAfter=3,
    )
    s_body = ParagraphStyle(
        "body",
        parent=styles["Normal"],
        fontSize=10,
        fontName="Helvetica",
        textColor=colors.HexColor("#374151"),
        leading=15,
        spaceAfter=4,
    )
    s_label = ParagraphStyle(
        "label",
        parent=styles["Normal"],
        fontSize=8,
        fontName="Helvetica-Bold",
        textColor=colors.white,
    )

    def agent_badge(label: str, color: str) -> Table:
        """Small colored badge showing which agent produced this section."""
        data = [[Paragraph(f"◉  {label}", s_label)]]
        t = Table(data, colWidths=[W])
        t.setStyle(
            TableStyle(
                [
                    ("BACKGROUND", (0, 0), (-1, -1), colors.HexColor(color)),
                    ("TOPPADDING", (0, 0), (-1, -1), 5),
                    ("BOTTOMPADDING", (0, 0), (-1, -1), 5),
                    ("LEFTPADDING", (0, 0), (-1, -1), 10),
                    ("ROUNDEDCORNERS", [4, 4, 4, 4]),
                ]
            )
        )
        return t

    def render_markdown_block(md_text: str) -> list:
        """Convert basic markdown to ReportLab flowables."""
        flowables = []
        for line in md_text.splitlines():
            line = line.strip()
            if not line:
                flowables.append(Spacer(1, 4))
                continue
            if line.startswith("## "):
                flowables.append(Paragraph(line[3:], s_section))
            elif line.startswith("### "):
                flowables.append(Paragraph(line[4:], s_h3))
            elif line.startswith("- ") or line.startswith("* "):
                clean = strip_markdown(line[2:])
                flowables.append(Paragraph(f"• {clean}", s_body))
            elif re.match(r"^\d+\.", line):
                clean = strip_markdown(line)
                flowables.append(Paragraph(clean, s_body))
            elif line.startswith("**") and line.endswith("**"):
                flowables.append(Paragraph(line[2:-2], s_h3))
            else:
                clean = strip_markdown(line)
                flowables.append(Paragraph(clean, s_body))
        return flowables

    story = []

    # ── Header
    story.append(Paragraph("Problem Analysis Report", s_title))
    story.append(
        Paragraph(
            f"{name}  ·  {role}  ·  {datetime.now().strftime('%d %B %Y')}",
            s_sub,
        )
    )
    story.append(HRFlowable(width=W, thickness=1.5, color=colors.HexColor("#2563eb"), spaceAfter=10))

    # ── Section 1: Problem Statement
    story.append(agent_badge("AGENT 1  —  Problem Analyst", "#1e3a5f"))
    story.append(Spacer(1, 6))
    story += render_markdown_block(analysis.problem_statement)

    # ── Section 2: Root Cause
    story.append(Spacer(1, 8))
    story.append(agent_badge("AGENT 2  —  Root Cause Analyst", "#1a4731"))
    story.append(Spacer(1, 6))
    story += render_markdown_block(analysis.root_causes)

    # ── Section 3: Solutions
    story.append(Spacer(1, 8))
    story.append(agent_badge("AGENT 3  —  Solution Brainstorm", "#4a1942"))
    story.append(Spacer(1, 6))
    story += render_markdown_block(analysis.solutions)

    # ── Section 4: Action Plan
    story.append(Spacer(1, 8))
    story.append(agent_badge("AGENT 4  —  Action Planner", "#7c2d12"))
    story.append(Spacer(1, 6))
    story += render_markdown_block(analysis.action_plan)

    # ── Section 5: Thinking Coach
    story.append(Spacer(1, 8))
    story.append(agent_badge("AGENT 5  —  Thinking Coach", "#312e81"))
    story.append(Spacer(1, 6))
    story += render_markdown_block(analysis.thinking_feedback)

    # ── Footer note
    story.append(Spacer(1, 16))
    story.append(HRFlowable(width=W, thickness=0.5, color=colors.HexColor("#e5e7eb"), spaceAfter=6))
    story.append(
        Paragraph(
            "Generated by the AI Intern Problem-Solving System  ·  Confidential",
            ParagraphStyle("footer", parent=styles["Normal"], fontSize=8, textColor=colors.HexColor("#9ca3af"), alignment=1),
        )
    )

    doc.build(story)
    buf.seek(0)
    return buf.read()


# ── API Endpoints ─────────────────────────────────────────────────────────────

@app.get("/")
def root():
    return {
        "service": "Intern Problem-Solving API",
        "version": "1.0.0",
        "agents": [
            "1. Problem Analyst",
            "2. Root Cause Analyst",
            "3. Solution Brainstorm",
            "4. Action Planner",
            "5. Thinking Coach",
        ],
        "endpoints": {
            "POST /analyze": "Run full 5-agent pipeline, returns JSON",
            "POST /analyze/stream": "Stream analysis as server-sent events",
            "POST /analyze/pdf": "Run pipeline + return downloadable PDF",
            "GET /health": "Health check",
        },
    }


@app.get("/health")
def health():
    return {"status": "ok", "timestamp": datetime.utcnow().isoformat()}


@app.post("/analyze", response_model=FullAnalysis)
def analyze(body: ProblemInput):
    """Run full 5-agent pipeline. Returns structured JSON."""
    if not body.content.strip():
        raise HTTPException(status_code=400, detail="Content cannot be empty.")

    content = decode_content(body.content)

    if len(content) < 30:
        raise HTTPException(status_code=400, detail="Content too short for meaningful analysis.")

    try:
        result = run_pipeline(
            content=content,
            name=body.intern_name or "Intern",
            role=body.intern_role or "AI Developer Intern",
            goal=body.intern_goal or "",
        )
        return result
    except anthropic.AuthenticationError:
        raise HTTPException(status_code=401, detail="Invalid Anthropic API key.")
    except anthropic.RateLimitError:
        raise HTTPException(status_code=429, detail="Rate limit reached. Please wait and retry.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/analyze/stream")
def analyze_stream(body: ProblemInput):
    """Stream each agent's output as server-sent events (SSE)."""
    if not body.content.strip():
        raise HTTPException(status_code=400, detail="Content cannot be empty.")

    # Decode PDF if needed before streaming starts
    resolved_content = decode_content(body.content)

    def event_stream():
        agents = [
            ("analyst", AGENT_ANALYST, "Problem Analyst"),
            ("root_cause", AGENT_ROOT_CAUSE, "Root Cause Analyst"),
            ("solutions", AGENT_SOLUTIONS, "Solution Brainstorm"),
            ("action_plan", AGENT_ACTION_PLANNER, "Action Planner"),
            ("thinking", AGENT_THINKING_COACH, "Thinking Coach"),
        ]

        context = {
            "content": resolved_content[:8000],
            "name": body.intern_name or "Intern",
            "role": body.intern_role or "AI Developer Intern",
            "goal": body.intern_goal or "",
        }
        accumulated = {}

        for key, system_prompt, label in agents:
            # Send agent start event
            yield f"data: {json.dumps({'event': 'agent_start', 'agent': key, 'label': label})}\n\n"

            # Build context-aware prompt for this agent
            if key == "analyst":
                user_msg = f"Intern: {context['name']} | Role: {context['role']} | Goal: {context['goal']}\n\nAnalyze this content:\n{context['content']}"
            elif key == "root_cause":
                user_msg = f"Problem:\n{accumulated.get('analyst','')}\n\nOriginal:\n{context['content'][:2000]}"
            elif key == "solutions":
                user_msg = f"Problem:\n{accumulated.get('analyst','')}\n\nRoot Causes:\n{accumulated.get('root_cause','')}"
            elif key == "action_plan":
                user_msg = f"Role: {context['role']}\n\nProblem:\n{accumulated.get('analyst','')}\n\nCauses:\n{accumulated.get('root_cause','')}\n\nSolutions:\n{accumulated.get('solutions','')}"
            else:
                user_msg = f"Original Input:\n{context['content'][:2500]}\n\nProblem:\n{accumulated.get('analyst','')}\n\nCauses:\n{accumulated.get('root_cause','')}"

            # Stream this agent's output
            agent_text = ""
            with client.messages.stream(
                model="claude-sonnet-4-6",
                max_tokens=800,
                system=system_prompt,
                messages=[{"role": "user", "content": user_msg}],
            ) as stream:
                for chunk in stream.text_stream:
                    agent_text += chunk
                    yield f"data: {json.dumps({'event': 'token', 'agent': key, 'text': chunk})}\n\n"

            accumulated[key] = agent_text
            yield f"data: {json.dumps({'event': 'agent_done', 'agent': key})}\n\n"

        yield f"data: {json.dumps({'event': 'done'})}\n\n"

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@app.post("/analyze/pdf")
def analyze_pdf(body: ProblemInput):
    """Run full pipeline and return a downloadable PDF report."""
    if not body.content.strip():
        raise HTTPException(status_code=400, detail="Content cannot be empty.")

    content = decode_content(body.content)

    try:
        analysis = run_pipeline(
            content=content,
            name=body.intern_name or "Intern",
            role=body.intern_role or "AI Developer Intern",
            goal=body.intern_goal or "",
        )
        pdf_bytes = build_pdf(
            analysis,
            name=body.intern_name or "Intern",
            role=body.intern_role or "AI Developer Intern",
        )
        filename = f"problem_analysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
        return StreamingResponse(
            io.BytesIO(pdf_bytes),
            media_type="application/pdf",
            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
        )
    except anthropic.AuthenticationError:
        raise HTTPException(status_code=401, detail="Invalid Anthropic API key.")
    except anthropic.RateLimitError:
        raise HTTPException(status_code=429, detail="Rate limit. Please wait and retry.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ── Dev Runner ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)