app.py

from dotenv import load_dotenv
load_dotenv()

import chainlit as cl
from memory.profile import load_profile, save_profile, profile_exists
from memory.prefs import load_prefs, save_prefs
from memory.vector_store import get_store, has_index
from graphs.profile_graph import profile_graph
from graphs.apply_graph import apply_graph
from nodes.draft import draft_email, draft_cover_letter
from nodes.extract import apply_profile_edit
from utils.scraper import extract_pdf_text

TONES = ["formal", "casual", "startup", "direct"]
_JOB_KEYWORDS = [
    "responsibilities", "requirements", "qualifications",
    "we are looking", "job description", "about the role",
    "what you'll do", "you will", "we offer", "apply now",
]
_VOCAB_PHRASES = {
    "rich": [
        "improve the vocab", "improve vocab", "better vocab", "richer vocab",
        "fancier words", "more advanced vocab", "elevate the language",
        "smarter vocab", "polish the vocab", "stronger vocabulary", "upgrade the vocab",
    ],
    "simple": [
        "simpler vocab", "simplify the vocab", "simplify the language", "lower the vocab",
        "easier words", "plainer language", "dumb it down", "simpler words",
        "tone down the vocab", "basic vocab", "casual words",
    ],
}


_PROFILE_EDIT_TRIGGERS = [
    "add my ", "add link", "add skill", "remove skill",
    "update my ", "edit my ", "change my ",
    "my email is", "my phone is", "my current role is",
    "update profile", "edit profile",
]
_PREF_EDIT_TRIGGERS = [
    "default tone", "tone preference", "set tone", "save tone",
    "change my preference", "update preference", "set default",
]


def _looks_like_profile_edit(lower: str) -> bool:
    return any(t in lower for t in _PROFILE_EDIT_TRIGGERS)


def _looks_like_pref_edit(lower: str) -> bool:
    return any(t in lower for t in _PREF_EDIT_TRIGGERS)


def _match_vocab(lower: str) -> str | None:
    for level, phrases in _VOCAB_PHRASES.items():
        if any(p in lower for p in phrases):
            return level
    return None


def _looks_like_job(text: str) -> bool:
    lower = text.lower()
    return sum(1 for kw in _JOB_KEYWORDS if kw in lower) >= 2


def _is_url(text: str) -> bool:
    return text.strip().startswith(("http://", "https://"))


def _pdf_text_from_elements(elements: list) -> str | None:
    for el in elements or []:
        mime = (getattr(el, "mime", "") or "").lower()
        name = (getattr(el, "name", "") or "").lower()
        path = getattr(el, "path", None)
        if path and (mime == "application/pdf" or name.endswith(".pdf")):
            pdf_text = extract_pdf_text(path)
            if pdf_text:
                return pdf_text
    return None


# ── startup ──────────────────────────────────────────────────────────────────

def _uid() -> str:
    return cl.user_session.get("user_id", "default")


@cl.on_chat_start
async def on_start():
    user_id = cl.context.session.id
    cl.user_session.set("user_id", user_id)

    if not profile_exists(user_id):
        cl.user_session.set("mode", "awaiting_resume")
        await cl.Message(
            "Hi! I'm **tailor-mail** — I draft tailored application emails and cover letters "
            "matched to your profile and the role you're targeting.\n\n"
            "To get started I need your resume. Paste the text below or upload a PDF and "
            "I'll extract your profile automatically."
        ).send()
        return

    profile = load_profile(user_id)
    prefs = load_prefs(user_id)

    if not has_index(user_id):
        await get_store(user_id).index_profile(profile)

    cl.user_session.set("mode", "ready")

    tone_note = f" Default tone: **{prefs['tone']}**." if prefs.get("tone") else ""
    await cl.Message(
        f"Welcome back, **{profile['name']}**!{tone_note}\n\n"
        "Paste a job URL or description and I'll draft your application.\n\n"
        "**Quick actions you can use any time:**\n"
        "- Say **\"cover letter\"** in your message to get a cover letter instead of an email\n"
        "- **\"add my GitHub: [url]\"** — add a link to your profile\n"
        "- **\"add [skill] to my skills\"** — update your skills\n"
        "- **\"set default tone to casual\"** — change your saved tone preference\n"
        "- `/profile` — view your full profile"
    ).send()


# ── messages ─────────────────────────────────────────────────────────────────

@cl.on_message
async def on_message(message: cl.Message):
    text = message.content.strip()
    mode = cl.user_session.get("mode", "ready")

    if mode == "awaiting_resume":
        await _handle_resume(_pdf_text_from_elements(message.elements) or text)
        return

    if mode == "awaiting_links":
        await _handle_links_reply(text)
        return

    if mode == "awaiting_manual_edit":
        await _handle_manual_edit_reply(text)
        return

    if text.startswith("/profile"):
        await _handle_profile_cmd()
        return

    output_type = "cover_letter" if "cover letter" in text.lower() else "email"

    if _is_url(text) or _looks_like_job(text):
        await _handle_job_input(text, output_type)
        return

    # Tone / vocabulary adjustment on an already-extracted job
    lower = text.lower()
    last_job = cl.user_session.get("last_job_facts")
    matched_tone = next((t for t in TONES if t in lower), None)
    matched_vocab = _match_vocab(lower)
    if last_job and (matched_tone or matched_vocab):
        tone = matched_tone or cl.user_session.get("last_tone") or load_prefs(_uid()).get("tone") or "formal"
        vocab_level = matched_vocab if matched_vocab else cl.user_session.get("vocab_level")
        extra_links = cl.user_session.get("extra_links")
        emphasis = cl.user_session.get("last_emphasis")
        await _draft_and_send(last_job, output_type, tone, save=bool(matched_tone), vocab_level=vocab_level, extra_links=extra_links, emphasis=emphasis)
        return

    if _looks_like_pref_edit(lower):
        await _handle_pref_edit(lower)
        return

    if _looks_like_profile_edit(lower):
        await _handle_profile_edit(text)
        return

    await cl.Message(
        "Paste a job posting URL or description and I'll draft your application.\n"
        "For a cover letter, say so in your message.\n"
        "You can also say things like 'add my GitHub link' or 'set default tone to casual'."
    ).send()


# ── actions ──────────────────────────────────────────────────────────────────

@cl.action_callback("copy_draft")
async def on_copy_draft(action: cl.Action):
    draft = cl.user_session.get("last_draft", "")
    if draft:
        await cl.Message(f"```\n{draft}\n```").send()


# ── handlers ─────────────────────────────────────────────────────────────────

async def _handle_resume(resume_text: str):
    if len(resume_text.split()) < 30:
        await cl.Message("That looks too short to be a resume. Paste the full text.").send()
        return

    thinking = await cl.Message("Extracting your profile...").send()

    try:
        result = await profile_graph.ainvoke({
            "resume_text": resume_text,
            "profile": None,
            "status": "",
            "user_id": _uid(),
        })
    except Exception:
        await cl.Message(
            "Something went wrong while extracting your profile. Please try pasting your resume again."
        ).send()
        return
    profile = result["profile"]

    social_note = ""
    gh = profile.get("social_data", {}).get("github", {})
    if gh.get("top_repos"):
        repos = ", ".join(r["name"] for r in gh["top_repos"][:3])
        social_note = f"\nGitHub repos pulled: **{repos}**"

    await cl.Message(
        f"Profile saved for **{profile['name']}** — {profile['current_role']}.\n"
        f"Skills: {', '.join(profile['skills'][:6])}"
        f"{social_note}"
    ).send()

    domain_res = await cl.AskActionMessage(
        content="What types of roles are you primarily targeting? This helps me tailor your drafts.",
        actions=[
            cl.Action(name="swe",     label="Software Engineering", payload={"domain": "software engineering"}),
            cl.Action(name="data",    label="Data & ML",            payload={"domain": "data science and machine learning"}),
            cl.Action(name="product", label="Product & Design",     payload={"domain": "product management and design"}),
            cl.Action(name="other",   label="Other / Mixed",        payload={"domain": None}),
        ],
        timeout=90,
    ).send()

    target_domain = domain_res.get("payload", {}).get("domain") if domain_res else None
    cl.user_session.set("target_domain", target_domain)
    cl.user_session.set("mode", "ready")

    domain_note = f" targeting **{target_domain}** roles" if target_domain else ""
    await cl.Message(
        f"All set{domain_note}! Paste a job URL or description whenever you're ready."
    ).send()


async def _handle_profile_cmd():
    profile = load_profile(_uid())
    if not profile:
        await cl.Message("No profile found. Paste your resume to set one up.").send()
        cl.user_session.set("mode", "awaiting_resume")
        return

    exp_lines = "\n".join(
        f"- {e['role']} at {e['company']}" for e in profile.get("experience", [])[:4]
    )
    await cl.Message(
        f"**{profile['name']}** — {profile['current_role']}\n\n"
        f"**Skills:** {', '.join(profile.get('skills', []))}\n\n"
        f"**Experience:**\n{exp_lines}\n\n"
        "To update your profile, paste a new resume."
    ).send()
    cl.user_session.set("mode", "awaiting_resume")


async def _handle_job_input(text: str, output_type: str):
    thinking = await cl.Message("Extracting job details...").send()

    try:
        result = await apply_graph.ainvoke({
            "raw_input": text,
            "job_text": None,
            "job_facts": None,
            "fetch_failed": False,
        })
    except Exception:
        await cl.Message(
            "Something went wrong while reading that job posting. "
            "Try pasting the job description text directly instead of the URL."
        ).send()
        return

    if result.get("fetch_failed"):
        await cl.Message(
            "I couldn't fetch that URL (the site may block automated requests) — "
            "for a more accurate draft, paste the job description text directly."
        ).send()

    job_facts = result["job_facts"]
    cl.user_session.set("last_job_facts", job_facts)
    cl.user_session.set("extra_links", None)
    cl.user_session.set("last_emphasis", None)

    label = "cover letter" if output_type == "cover_letter" else "email"

    angle_res = await cl.AskActionMessage(
        content=f"Got it — **{job_facts['role']}** at **{job_facts['company']}**.\n\nAnything to emphasise in this {label}?",
        actions=[
            cl.Action(name="technical",  label="Technical depth",      payload={"emphasis": "technical expertise and engineering depth"}),
            cl.Action(name="leadership", label="Leadership",           payload={"emphasis": "leadership and team management"}),
            cl.Action(name="projects",   label="Projects & portfolio", payload={"emphasis": "personal projects and open-source work"}),
            cl.Action(name="domain",     label="Domain expertise",     payload={"emphasis": "domain and industry knowledge"}),
            cl.Action(name="skip",       label="No specific focus",    payload={"emphasis": None}),
        ],
        timeout=90,
    ).send()

    emphasis = angle_res.get("payload", {}).get("emphasis") if angle_res else None
    cl.user_session.set("last_emphasis", emphasis)

    prefs = load_prefs(_uid())
    saved_tone = prefs.get("tone")
    vocab_level = cl.user_session.get("vocab_level")

    if saved_tone:
        await cl.Message(f"Drafting your {label} ({saved_tone} tone)...").send()
        await _draft_and_send(job_facts, output_type, saved_tone, save=False, vocab_level=vocab_level, emphasis=emphasis)
    else:
        res = await cl.AskActionMessage(
            content=f"What tone for the {label}?",
            actions=[
                cl.Action(name=t, label=t.capitalize(), payload={"tone": t})
                for t in TONES
            ],
        ).send()
        tone = res.get("payload", {}).get("tone", "formal") if res else "formal"
        await _draft_and_send(job_facts, output_type, tone, save=True, vocab_level=vocab_level, emphasis=emphasis)


async def _draft_and_send(
    job_facts: dict,
    output_type: str,
    tone: str,
    save: bool,
    vocab_level: str | None = None,
    extra_links: list | None = None,
    emphasis: str | None = None,
):
    profile = load_profile(_uid())
    cl.user_session.set("last_tone", tone)
    cl.user_session.set("vocab_level", vocab_level)
    if extra_links is not None:
        cl.user_session.set("extra_links", extra_links)
    else:
        extra_links = cl.user_session.get("extra_links")
    if emphasis is not None:
        cl.user_session.set("last_emphasis", emphasis)
    else:
        emphasis = cl.user_session.get("last_emphasis")

    msg = cl.Message(
        content="",
        actions=[cl.Action(name="copy_draft", label="Copy", payload={})],
    )
    await msg.send()

    gen = (
        draft_email(profile, job_facts, tone, vocab_level, extra_links, _uid(), emphasis)
        if output_type == "email"
        else draft_cover_letter(profile, job_facts, tone, vocab_level, extra_links, _uid(), emphasis)
    )
    tokens = []
    try:
        async for chunk in gen:
            await msg.stream_token(chunk)
            tokens.append(chunk)
    except Exception:
        await msg.update()
        await cl.Message(
            "Something went wrong while drafting — please try again."
        ).send()
        return
    await msg.update()
    cl.user_session.set("last_draft", "".join(tokens))

    if save:
        save_prefs({"tone": tone}, _uid())
        await cl.Message(f"Tone saved as **{tone}** for next time.").send()

    await _ask_edit_choices(job_facts, output_type, tone, vocab_level, extra_links, emphasis)


async def _handle_pref_edit(lower: str):
    matched_tone = next((t for t in TONES if t in lower), None)
    if matched_tone:
        save_prefs({"tone": matched_tone}, _uid())
        await cl.Message(f"Default tone updated to **{matched_tone}**.").send()
    else:
        prefs = load_prefs(_uid())
        current = prefs.get("tone") or "not set"
        await cl.Message(
            f"Current preferences: tone = **{current}**.\n\n"
            f"Say 'set default tone to formal / casual / startup / direct' to change it."
        ).send()


async def _handle_profile_edit(text: str):
    profile = load_profile(_uid())
    if not profile:
        await cl.Message("No profile found. Paste your resume to set one up.").send()
        cl.user_session.set("mode", "awaiting_resume")
        return

    thinking = await cl.Message("Updating your profile...").send()
    try:
        edit = await apply_profile_edit(profile, text)
    except Exception:
        await cl.Message(
            "Something went wrong while updating your profile. Please try again."
        ).send()
        return

    changed = []
    if edit.name:
        profile["name"] = edit.name
        changed.append(f"name → **{edit.name}**")
    if edit.email:
        profile["email"] = edit.email
        changed.append(f"email → **{edit.email}**")
    if edit.phone:
        profile["phone"] = edit.phone
        changed.append(f"phone → **{edit.phone}**")
    if edit.current_role:
        profile["current_role"] = edit.current_role
        changed.append(f"current role → **{edit.current_role}**")
    if edit.summary:
        profile["summary"] = edit.summary
        changed.append("summary updated")
    if edit.skills_add:
        profile["skills"] = list(dict.fromkeys(profile.get("skills", []) + edit.skills_add))
        changed.append(f"skills added: **{', '.join(edit.skills_add)}**")
    if edit.skills_remove:
        removed = [s for s in edit.skills_remove if s in profile.get("skills", [])]
        profile["skills"] = [s for s in profile.get("skills", []) if s not in edit.skills_remove]
        if removed:
            changed.append(f"skills removed: **{', '.join(removed)}**")
    if edit.social_links_add:
        profile["social_links"] = list(dict.fromkeys(
            profile.get("social_links", []) + edit.social_links_add
        ))
        changed.append(f"links added: **{', '.join(edit.social_links_add)}**")
    if edit.social_links_remove:
        profile["social_links"] = [
            l for l in profile.get("social_links", []) if l not in edit.social_links_remove
        ]
        changed.append("links removed")
    if edit.education_add:
        profile["education"] = profile.get("education", []) + edit.education_add
        changed.append(f"education added: **{', '.join(edit.education_add)}**")

    if not changed:
        await cl.Message(
            "I couldn't figure out what to update. Try being more specific — "
            "e.g. 'add my GitHub: https://github.com/you' or 'add React to my skills'."
        ).send()
        return

    save_profile(profile, _uid())
    await get_store(_uid()).index_profile(profile)

    await cl.Message(
        "Profile updated:\n" + "\n".join(f"- {c}" for c in changed)
    ).send()


async def _ask_edit_choices(
    job_facts: dict,
    output_type: str,
    tone: str,
    vocab_level: str | None,
    extra_links: list | None,
    emphasis: str | None,
):
    other_type = "cover letter" if output_type == "email" else "email"
    label = "email" if output_type == "email" else "cover letter"

    res = await cl.AskActionMessage(
        content=f"Want to refine this {label}?",
        actions=[
            cl.Action(name="change_tone",  label="Change tone",            payload={"action": "change_tone"}),
            cl.Action(name="change_angle", label="Change emphasis",         payload={"action": "change_angle"}),
            cl.Action(name="add_links",    label="Add / edit links",        payload={"action": "add_links"}),
            cl.Action(name="manual_edit",  label="Edit manually",           payload={"action": "manual_edit"}),
            cl.Action(name="switch",       label=f"Switch to {other_type}", payload={"action": "switch"}),
            cl.Action(name="done",         label="Done",                    payload={"action": "done"}),
        ],
        timeout=120,
    ).send()

    if not res:
        return

    action = res.get("payload", {}).get("action")

    if not action or action == "done":
        return
    elif action == "change_tone":
        tone_res = await cl.AskActionMessage(
            content="Which tone?",
            actions=[cl.Action(name=t, label=t.capitalize(), payload={"tone": t}) for t in TONES],
        ).send()
        new_tone = tone_res.get("payload", {}).get("tone", tone) if tone_res else tone
        await _draft_and_send(job_facts, output_type, new_tone, save=True, vocab_level=vocab_level, extra_links=extra_links, emphasis=emphasis)
    elif action == "change_angle":
        await _handle_change_angle(job_facts, output_type, tone, vocab_level, extra_links)
    elif action == "switch":
        new_type = "cover_letter" if output_type == "email" else "email"
        await _draft_and_send(job_facts, new_type, tone, save=False, vocab_level=vocab_level, extra_links=extra_links, emphasis=emphasis)
    elif action == "add_links":
        await _handle_add_links(job_facts, output_type, tone, vocab_level, extra_links, emphasis)
    elif action == "manual_edit":
        cl.user_session.set("mode", "awaiting_manual_edit")
        cl.user_session.set("pending_draft", {
            "job_facts": job_facts, "output_type": output_type, "tone": tone,
            "vocab_level": vocab_level, "extra_links": extra_links, "emphasis": emphasis,
        })
        draft = cl.user_session.get("last_draft", "")
        await cl.Message(
            f"Here's the current draft — edit it however you like and paste your version below:\n\n"
            f"```\n{draft}\n```"
        ).send()


async def _handle_change_angle(
    job_facts: dict,
    output_type: str,
    tone: str,
    vocab_level: str | None,
    extra_links: list | None,
):
    current = cl.user_session.get("last_emphasis")
    hint = f"\n\nCurrent: _{current}_" if current else ""

    res = await cl.AskActionMessage(
        content=f"What should I emphasise?{hint}",
        actions=[
            cl.Action(name="technical",  label="Technical depth",      payload={"emphasis": "technical expertise and engineering depth"}),
            cl.Action(name="leadership", label="Leadership",           payload={"emphasis": "leadership and team management"}),
            cl.Action(name="projects",   label="Projects & portfolio", payload={"emphasis": "personal projects and open-source work"}),
            cl.Action(name="domain",     label="Domain expertise",     payload={"emphasis": "domain and industry knowledge"}),
            cl.Action(name="clear",      label="Clear / no focus",     payload={"emphasis": None}),
        ],
        timeout=120,
    ).send()

    if not res:
        return

    emphasis = res.get("payload", {}).get("emphasis")
    cl.user_session.set("last_emphasis", emphasis)

    await _draft_and_send(job_facts, output_type, tone, save=False, vocab_level=vocab_level, extra_links=extra_links, emphasis=emphasis)


async def _handle_add_links(
    job_facts: dict,
    output_type: str,
    tone: str,
    vocab_level: str | None,
    current_extra: list | None,
    emphasis: str | None = None,
):
    profile = load_profile(_uid()) or {}
    existing = (profile.get("social_links") or []) + (current_extra or [])
    hint = f"\n\nCurrently included: {', '.join(existing)}" if existing else ""

    cl.user_session.set("mode", "awaiting_links")
    cl.user_session.set("pending_draft", {
        "job_facts": job_facts,
        "output_type": output_type,
        "tone": tone,
        "vocab_level": vocab_level,
        "current_extra": current_extra,
        "emphasis": emphasis,
    })

    await cl.Message(
        f"Paste the links to include — one per line or comma-separated "
        f"(GitHub, LinkedIn, LeetCode, portfolio…).{hint}"
    ).send()


async def _handle_manual_edit_reply(text: str):
    pending = cl.user_session.get("pending_draft", {})
    cl.user_session.set("mode", "ready")
    cl.user_session.set("last_draft", text)

    await cl.Message(f"**Your edited version:**\n\n{text}").send()
    await _ask_edit_choices(
        pending["job_facts"], pending["output_type"], pending["tone"],
        pending.get("vocab_level"), pending.get("extra_links"), pending.get("emphasis"),
    )


async def _handle_links_reply(text: str):
    pending = cl.user_session.get("pending_draft", {})
    cl.user_session.set("mode", "ready")

    new_links = [l.strip() for l in text.replace(",", "\n").splitlines() if l.strip()]
    if not new_links:
        await cl.Message("No links found — paste them one per line or comma-separated.").send()
        return

    current_extra = pending.get("current_extra") or []
    merged = list(dict.fromkeys(current_extra + new_links))
    cl.user_session.set("extra_links", merged)

    await _draft_and_send(
        pending["job_facts"], pending["output_type"], pending["tone"],
        save=False,
        vocab_level=pending.get("vocab_level"),
        extra_links=merged,
        emphasis=pending.get("emphasis"),
    )