!pip install -q langgraph langchain-openai pydantic requests

import os
import json
import requests
from typing import TypedDict, Optional
from datetime import datetime, timedelta, timezone
from uuid import uuid4

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")      # required for generation/review
PUBLORA_API_KEY = os.environ.get("PUBLORA_API_KEY")    # optional -> dry-run if unset
PUBLORA_BASE = "https://api.publora.com/api/v1"

DRY_RUN = True   # True -> create drafts (or print). False -> actually schedule at schedule_time.
MAX_ITERS = 2    # max self-revision rounds before publishing what we have

PLATFORM_SPECS = {
    "twitter":  {"max_chars": 280,  "style": "punchy and concrete; a strong hook in the first line; at most 1-2 hashtags"},
    "linkedin": {"max_chars": 3000, "style": "professional and value-driven; short paragraphs; one clear takeaway; no hashtag spam"},
    "bluesky":  {"max_chars": 300,  "style": "conversational and authentic; minimal hashtags; a builder-to-builder tone"},
    "threads":  {"max_chars": 500,  "style": "casual and engaging; invites replies; light, human voice"},
    "mastodon": {"max_chars": 500,  "style": "informative and community-friendly; plain, non-salesy"},
}

def publora_request(method: str, path: str, body: dict | None = None, idempotency_key: str | None = None) -> dict:
    """Authenticated call to the Publora REST API."""
    headers = {"x-publora-key": PUBLORA_API_KEY}
    if idempotency_key:
        headers["Idempotency-Key"] = idempotency_key
    resp = requests.request(method, f"{PUBLORA_BASE}/{path}", headers=headers, json=body, timeout=30)
    if resp.status_code == 401:
        raise RuntimeError("Invalid PUBLORA_API_KEY - generate one at publora.com -> API")
    resp.raise_for_status()
    return resp.json()


def get_connections() -> list:
    """Return the list of connected social accounts for this key."""
    return publora_request("GET", "platform-connections")["connections"]


def create_post(content: str, platform_id: str, scheduled_time: str | None = None) -> dict:
    """Schedule (or draft) one post. One platform_id per call so each can carry tailored content."""
    body = {"content": content, "platforms": [platform_id]}
    if scheduled_time:
        body["scheduledTime"] = scheduled_time
    return publora_request("POST", "create-post", body=body, idempotency_key=str(uuid4()))

def get_llm():
    if not OPENAI_API_KEY:
        raise RuntimeError("Set OPENAI_API_KEY to generate content.")
    return ChatOpenAI(model="gpt-4o-mini", temperature=0.7, api_key=OPENAI_API_KEY)


class PlatformPost(BaseModel):
    content: str = Field(description="The ready-to-publish post text for this platform.")


class Critique(BaseModel):
    approved: bool = Field(description="True if the post is ready to publish as-is.")
    note: str = Field(description="If not approved, a short, specific instruction on how to fix it.")

class AgentState(TypedDict):
    idea: str                      # the raw input idea
    target_platforms: list         # platforms the user asked for
    connections: dict              # platform -> real Publora platformId
    drafts: dict                   # platform -> generated content
    feedback: dict                 # platform -> revision instruction (from review)
    iteration: int                 # how many review rounds have run
    schedule_time: Optional[str]   # ISO 8601 UTC; used only when DRY_RUN is False
    results: list                  # per-platform publish results

def fetch_connections(state: AgentState) -> dict:
    targets = state["target_platforms"]
    if PUBLORA_API_KEY:
        mapping = {}
        for c in get_connections():
            prefix = c["platformId"].split("-", 1)[0]      # "twitter-123" -> "twitter"
            mapping.setdefault(prefix, c["platformId"])     # first account per platform
        connections = {p: mapping[p] for p in targets if p in mapping}
        missing = [p for p in targets if p not in mapping]
        if missing:
            print(f"Not connected in your Publora account, skipping: {missing}")
    else:
        connections = {p: f"{p}-DEMO123" for p in targets}
        print("No PUBLORA_API_KEY set - using mock connections (dry-run).")
    print(f"Publishing to: {list(connections)}")
    return {"connections": connections}

def generate(state: AgentState) -> dict:
    llm = get_llm().with_structured_output(PlatformPost)
    drafts = dict(state.get("drafts") or {})
    feedback = state.get("feedback") or {}
    platforms = [p for p in state["target_platforms"] if p in state["connections"]]

    for p in platforms:
        if p not in PLATFORM_SPECS:
            print(f"No spec for '{p}' - add it to PLATFORM_SPECS. Skipping.")
            continue
        if p in drafts and p not in feedback:
            continue  # already approved, leave it
        spec = PLATFORM_SPECS[p]
        instruction = (
            f"You are a social media copywriter. Write a single {p} post about the idea below.\n"
            f"Idea: {state['idea']}\n\n"
            f"Constraints:\n"
            f"- Hard limit: {spec['max_chars']} characters.\n"
            f"- Voice: {spec['style']}.\n"
            f"- No preamble, no quotes around the post, return only the post text."
        )
        if p in feedback:
            instruction += f"\n\nRevise your previous draft. Reviewer feedback: {feedback[p]}\nPrevious draft: {drafts.get(p, '')}"
        result = llm.invoke(instruction)
        drafts[p] = result.content.strip()
    return {"drafts": drafts, "feedback": {}}

def review(state: AgentState) -> dict:
    llm = get_llm().with_structured_output(Critique)
    feedback = {}
    for p, content in state["drafts"].items():
        spec = PLATFORM_SPECS[p]
        if len(content) > spec["max_chars"]:
            feedback[p] = f"Too long: {len(content)}/{spec['max_chars']} chars. Cut it down while keeping the hook."
            continue
        verdict = llm.invoke(
            f"Review this {p} post. Platform voice: {spec['style']}. "
            f"Approve only if it is on-idea, on-voice, and publish-ready.\n\nPost:\n{content}"
        )
        if not verdict.approved:
            feedback[p] = verdict.note
    kept = {p: state["drafts"][p] for p in state["drafts"] if p not in feedback}
    print(f"Review round {state.get('iteration', 0) + 1}: {len(kept)} approved, {len(feedback)} need revision.")
    return {"feedback": feedback, "iteration": state.get("iteration", 0) + 1}

def route_after_review(state: AgentState) -> str:
    if state["feedback"] and state["iteration"] < MAX_ITERS:
        return "generate"
    return "publish"

def publish(state: AgentState) -> dict:
    results = []
    for p, content in state["drafts"].items():
        platform_id = state["connections"][p]
        if not PUBLORA_API_KEY:
            print(f"[DRY-RUN] would POST create-post -> {platform_id}: {content[:60]}...")
            results.append({"platform": p, "status": "dry-run", "postGroupId": f"dry-{uuid4().hex[:8]}"})
            continue
        scheduled = None if DRY_RUN else state.get("schedule_time")
        try:
            resp = create_post(content, platform_id, scheduled_time=scheduled)
            results.append({
                "platform": p,
                "status": "scheduled" if scheduled else "draft",
                "postGroupId": resp.get("postGroupId"),
                "scheduledTime": resp.get("scheduledTime"),
            })
        except Exception as exc:  # isolate one platform's failure from the rest
            print(f"Failed to publish to {p}: {exc}")
            results.append({"platform": p, "status": "error", "error": str(exc)})
    return {"results": results}

builder = StateGraph(AgentState)
builder.add_node("fetch_connections", fetch_connections)
builder.add_node("generate", generate)
builder.add_node("review", review)
builder.add_node("publish", publish)

builder.set_entry_point("fetch_connections")
builder.add_edge("fetch_connections", "generate")
builder.add_edge("generate", "review")
builder.add_conditional_edges("review", route_after_review, {"generate": "generate", "publish": "publish"})
builder.add_edge("publish", END)

agent = builder.compile()

initial_state = {
    "idea": "We just shipped a REST + MCP API that lets AI agents publish to 10 social platforms with a single call.",
    "target_platforms": ["twitter", "linkedin", "bluesky"],
    "connections": {},
    "drafts": {},
    "feedback": {},
    "iteration": 0,
    "schedule_time": (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat(),
    "results": [],
}

final_state = agent.invoke(initial_state)

print("\n===== DRAFTS =====")
for platform, content in final_state["drafts"].items():
    print(f"\n--- {platform} ({len(content)} chars) ---\n{content}")

print("\n===== PUBLISH RESULTS =====")
print(json.dumps(final_state["results"], indent=2))