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 jinja2 import Template
from pydantic import BaseModel
from weasyprint import HTML

# ── App Setup ─────────────────────────────────────────────────────────────────
app = FastAPI(
    title="Intern Problem-Solving API",
    description="5-step pipeline: Analysis → Root Cause → Solutions → Action Plan → Reflection",
    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
    user_name: Optional[str] = "Anonymous"
    user_role: Optional[str] = ""
    user_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 person 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 person can start immediately. No vague advice."""


AGENT_THINKING_COACH = """You are the Thinking Coach Agent.
You are an encouraging but honest coach helping the person think more clearly.
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 for this person.

### For Next Time
3 specific things to do differently next time this type of problem comes up.

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"""
Name: {name}
Role: {role if role else "Not specified"}
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"""Role: {role}\nGoal: {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:\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 ─────────────────────────────────────────────────────────────

TEMPLATE_PATH = os.path.join(os.path.dirname(__file__), "report_template.html")


def md_to_html(text: str) -> str:
    """Convert basic markdown to clean HTML for the PDF template."""
    lines = text.splitlines()
    html_parts = []
    in_ul = False

    for line in lines:
        line = line.strip()

        # Close open list if needed
        if in_ul and not (line.startswith("- ") or line.startswith("* ")):
            html_parts.append("</ul>")
            in_ul = False

        if not line:
            continue
        elif line.startswith("## "):
            content = re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", line[3:])
            html_parts.append(f"<h2>{content}</h2>")
        elif line.startswith("### "):
            content = re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", line[4:])
            html_parts.append(f"<h3>{content}</h3>")
        elif line.startswith("- ") or line.startswith("* "):
            if not in_ul:
                html_parts.append("<ul>")
                in_ul = True
            content = re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", line[2:])
            html_parts.append(f"<li>{content}</li>")
        elif re.match(r"^\d+\.", line):
            content = re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", line)
            html_parts.append(f"<p>{content}</p>")
        else:
            content = re.sub(r"\*\*(.*?)\*\*", r"<strong>\1</strong>", line)
            html_parts.append(f"<p>{content}</p>")

    if in_ul:
        html_parts.append("</ul>")

    return "\n".join(html_parts)


def build_pdf(analysis: FullAnalysis, name: str, role: str) -> bytes:
    """Render Jinja2 HTML template and convert to PDF via WeasyPrint."""
    with open(TEMPLATE_PATH, "r", encoding="utf-8") as f:
        template = Template(f.read())

    sections = [
        {"title": "Problem Analysis", "html": md_to_html(analysis.problem_statement)},
        {"title": "Root Cause",        "html": md_to_html(analysis.root_causes)},
        {"title": "Solutions",         "html": md_to_html(analysis.solutions)},
        {"title": "Action Plan",       "html": md_to_html(analysis.action_plan)},
        {"title": "Reflection",        "html": md_to_html(analysis.thinking_feedback)},
    ]

    html_str = template.render(
        user_name=name,
        date=datetime.now().strftime("%d %B %Y"),
        sections=sections,
        page_breaks={3, 4},   # start Solutions and Action Plan on new page
    )

    pdf_bytes = HTML(string=html_str, base_url=None).write_pdf()
    return pdf_bytes


# ── 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.user_name or "Anonymous",
            role=body.user_role or "",
            goal=body.user_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 Analysis"),
            ("root_cause", AGENT_ROOT_CAUSE, "Root Cause"),
            ("solutions", AGENT_SOLUTIONS, "Solutions"),
            ("action_plan", AGENT_ACTION_PLANNER, "Action Plan"),
            ("thinking", AGENT_THINKING_COACH, "Reflection"),
        ]

        context = {
            "content": resolved_content[:8000],
            "name": body.user_name or "Anonymous",
            "role": body.user_role or "",
            "goal": body.user_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"Name: {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.user_name or "Anonymous",
            role=body.user_role or "",
            goal=body.user_goal or "",
        )
        pdf_bytes = build_pdf(
            analysis,
            name=body.user_name or "Anonymous",
            role=body.user_role or "",
        )
        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)