from __future__ import annotations

import json
from typing import TypedDict, List, Dict, Any, Literal, Optional
from datetime import datetime, timezone

from langgraph.graph import StateGraph, END

from .llm import worker_generate, supervisor_generate


class SupervisorVerdict(TypedDict):
    approved: bool
    issues: List[str]
    suggestions: List[str]


class EditorialState(TypedDict, total=False):
    # Inputs
    intro_text: str
    blog_links: List[str]

    # Working outputs
    subject_lines: List[str]
    newsletter_md: str

    # Supervision loop
    revision: int
    max_revisions: int
    verdict: SupervisorVerdict
    history: List[Dict[str, Any]]  # store per-revision summaries

    # Final report
    report: Dict[str, Any]


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def node_subject_lines(state: EditorialState) -> EditorialState:
    intro = state["intro_text"]
    links = state["blog_links"]

    prompt = (
        "Create 6 concise newsletter subject lines.\n"
        "Mix calm, energetic, and curious tones.\n"
        "Rules:\n"
        "- No clickbait\n"
        "- No emojis in subject lines\n"
        "- Keep each under 80 characters\n\n"
        f"Intro:\n{intro}\n\n"
        "Links:\n" + "\n".join(f"- {u}" for u in links)
    )

    text = worker_generate(prompt)
    lines = [l.strip("- ").strip() for l in text.splitlines() if l.strip()]
    subjects = [s for s in lines if 6 <= len(s) <= 80][:6] or ["AI notes + links (draft)"]

    state["subject_lines"] = subjects
    return state


def node_draft(state: EditorialState) -> EditorialState:
    intro = state["intro_text"]
    links = state["blog_links"]
    subjects = state.get("subject_lines", [])

    prompt = (
        "Write a Markdown newsletter draft.\n"
        "Audience: developers interested in AI and Python.\n"
        "Tone: friendly, clear, practical.\n"
        "Structure:\n"
        "1) Short greeting\n"
        "2) Intro paragraph (use the provided intro)\n"
        "3) A bullet list of links with 1–2 sentence summaries each\n"
        "4) Short closing\n\n"
        f"Intro:\n{intro}\n\n"
        "Links:\n" + "\n".join(f"- {u}" for u in links) + "\n\n"
        "Candidate subjects:\n" + "\n".join(f"- {s}" for s in subjects)
    )

    state["newsletter_md"] = worker_generate(prompt)
    return state


def _safe_parse_verdict(text: str) -> SupervisorVerdict:
    """
    Supervisor must return strict JSON, but we still guard against weird output.
    """
    try:
        data = json.loads(text)
        approved = bool(data.get("approved", False))
        issues = data.get("issues") or []
        suggestions = data.get("suggestions") or []
        if not isinstance(issues, list):
            issues = [str(issues)]
        if not isinstance(suggestions, list):
            suggestions = [str(suggestions)]
        return {"approved": approved, "issues": [str(x) for x in issues], "suggestions": [str(x) for x in suggestions]}
    except Exception:
        return {
            "approved": False,
            "issues": ["Supervisor returned non-JSON output."],
            "suggestions": ["Return strict JSON only: {approved, issues[], suggestions[]}."],
        }


def node_supervisor_check(state: EditorialState) -> EditorialState:
    intro = state["intro_text"]
    links = state["blog_links"]
    draft = state.get("newsletter_md", "")
    revision = state.get("revision", 0)

    prompt = (
        "You are a strict QA supervisor for a newsletter draft.\n"
        "Return ONLY valid JSON with keys: approved (bool), issues (array of strings), suggestions (array of strings).\n\n"
        "Checklist:\n"
        "- Draft is in Markdown.\n"
        "- Contains the intro content (or clearly incorporates it).\n"
        "- Includes ALL provided links exactly (no hallucinated links).\n"
        "- Each link has a 1–2 sentence summary.\n"
        "- Has a greeting and a closing.\n"
        "- Tone is friendly, clear, practical.\n\n"
        f"Provided intro:\n{intro}\n\n"
        "Provided links:\n" + "\n".join(f"- {u}" for u in links) + "\n\n"
        f"Draft (revision {revision}):\n{draft}\n\n"
        "Now return JSON only."
    )

    text = supervisor_generate(prompt)

    # If OPENAI_API_KEY isn't set, supervisor_generate returns None.
    # For this post, we want supervision enabled. But we still handle it gracefully.
    if text is None:
        verdict: SupervisorVerdict = {
            "approved": True,
            "issues": ["Supervisor disabled (no OPENAI_API_KEY)."],
            "suggestions": [],
        }
    else:
        verdict = _safe_parse_verdict(text)

    state["verdict"] = verdict

    # Append history snapshot
    history = state.get("history", [])
    history.append(
        {
            "ts": _utc_now(),
            "revision": revision,
            "approved": verdict["approved"],
            "issue_count": len(verdict["issues"]),
            "issues": verdict["issues"][:10],
        }
    )
    state["history"] = history

    return state


def node_revise(state: EditorialState) -> EditorialState:
    """
    Ask the worker model to revise the draft based on supervisor issues.
    """
    verdict = state.get("verdict", {"approved": False, "issues": [], "suggestions": []})
    issues = verdict.get("issues", [])
    suggestions = verdict.get("suggestions", [])
    draft = state.get("newsletter_md", "")
    links = state.get("blog_links", [])
    intro = state.get("intro_text", "")

    prompt = (
        "Revise the newsletter draft in Markdown.\n"
        "You MUST address the supervisor issues.\n"
        "Rules:\n"
        "- Keep it friendly, clear, practical.\n"
        "- Include all provided links exactly.\n"
        "- Each link gets 1–2 sentence summary.\n"
        "- Keep structure: greeting, intro, links list, closing.\n\n"
        f"Intro:\n{intro}\n\n"
        "Links:\n" + "\n".join(f"- {u}" for u in links) + "\n\n"
        "Supervisor issues:\n" + "\n".join(f"- {x}" for x in issues) + "\n\n"
        "Supervisor suggestions:\n" + "\n".join(f"- {x}" for x in suggestions) + "\n\n"
        "Current draft:\n" + draft
    )

    state["newsletter_md"] = worker_generate(prompt)
    return state


def route_after_supervisor(state: EditorialState) -> Literal["finalize", "revise", "give_up"]:
    verdict = state.get("verdict", {"approved": False, "issues": [], "suggestions": []})
    approved = bool(verdict.get("approved", False))

    revision = int(state.get("revision", 0))
    max_revisions = int(state.get("max_revisions", 2))

    if approved:
        return "finalize"

    # Not approved → can we retry?
    if revision < max_revisions:
        return "revise"

    return "give_up"


def node_bump_revision(state: EditorialState) -> EditorialState:
    state["revision"] = int(state.get("revision", 0)) + 1
    return state


def node_finalize_report(state: EditorialState) -> EditorialState:
    verdict = state.get("verdict", {"approved": False, "issues": [], "suggestions": []})

    state["report"] = {
        "created_at": _utc_now(),
        "approved": bool(verdict.get("approved", False)),
        "final_revision": int(state.get("revision", 0)),
        "max_revisions": int(state.get("max_revisions", 2)),
        "link_count": len(state.get("blog_links", [])),
        "subject_count": len(state.get("subject_lines", [])),
        "newsletter_chars": len(state.get("newsletter_md", "")),
        "history": state.get("history", []),
        "final_issues": verdict.get("issues", []),
        "notes": [
            "Draft only. No sending performed.",
            "Next steps: tool boundaries (MCP), notifications, and safe actions behind approvals.",
        ],
    }
    return state


def build_graph():
    g = StateGraph(EditorialState)

    g.add_node("subjects", node_subject_lines)
    g.add_node("draft", node_draft)
    g.add_node("supervisor", node_supervisor_check)
    g.add_node("revise", node_revise)
    g.add_node("bump_revision", node_bump_revision)
    g.add_node("finalize", node_finalize_report)

    g.set_entry_point("subjects")
    g.add_edge("subjects", "draft")
    g.add_edge("draft", "supervisor")

    # Conditional routing after supervisor
    g.add_conditional_edges(
        "supervisor",
        route_after_supervisor,
        {
            "finalize": "finalize",
            "revise": "revise",
            "give_up": "finalize",  # finalize report even when giving up
        },
    )

    # If we revise, bump revision, then check again
    g.add_edge("revise", "bump_revision")
    g.add_edge("bump_revision", "supervisor")

    g.add_edge("finalize", END)

    return g.compile()

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import List

from .graph import build_graph


def parse_links(raw: str) -> List[str]:
    items = [x.strip() for x in raw.split(",") if x.strip()]
    # de-dup keep order
    out = []
    seen = set()
    for u in items:
        if u not in seen:
            out.append(u)
            seen.add(u)
    return out


def main():
    parser = argparse.ArgumentParser(description="LangGraph newsletter drafter with supervisor loop")
    parser.add_argument("--intro", required=True, help="Intro paragraph for the newsletter")
    parser.add_argument("--links", required=True, help="Comma-separated list of URLs")
    parser.add_argument("--out", default="out", help="Output folder")
    args = parser.parse_args()

    max_revisions = int(os.getenv("MAX_REVISIONS", "2"))

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    graph = build_graph()

    state = {
        "intro_text": args.intro,
        "blog_links": parse_links(args.links),
        "revision": 0,
        "max_revisions": max_revisions,
        "history": [],
    }

    result = graph.invoke(state)

    (out_dir / "newsletter.md").write_text(result.get("newsletter_md", ""), encoding="utf-8")
    (out_dir / "subject_lines.txt").write_text("\n".join(result.get("subject_lines", [])) + "\n", encoding="utf-8")
    (out_dir / "report.json").write_text(json.dumps(result.get("report", {}), indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    print(f"Done. Wrote outputs to: {out_dir.resolve()}")
    print(" - newsletter.md")
    print(" - subject_lines.txt")
    print(" - report.json")


if __name__ == "__main__":
    main()