Thrive Career Wellness:
Automated Outplacement Platform

"""
Data schemas for the Thrive Wellness Career Platform.

Defines Pydantic models for candidate profiles, AI agent strategies, 
and job matching analytics to ensure strict type validation across the system.
"""

from pydantic import BaseModel, Field
from typing import List, Optional, Union


# --- 1. CANDIDATE PROFILE COMPONENTS ---

class Experience(BaseModel):
    """Represents a single professional role in a candidate's work history."""
    company: str
    location: str
    role: str
    duration: str
    achievements: List[str]


class Project(BaseModel):
    """Represents a technical or professional project in a candidate's portfolio."""
    title: str
    tech_stack: str
    date: str
    description_bullets: List[str]


class Education(BaseModel):
    """Represents a formal academic degree or certification."""
    school: str
    location: str
    degree: str
    year: str


class CandidateProfile(BaseModel):
    """The central data model for a candidate's full professional identity."""
    full_name: str
    contact_email: str
    phone: str
    linkedin: str
    summary: str
    skills: List[str]
    experience_history: List[Experience]
    projects: List[Project]
    education: List[Education]
    target_role: Optional[str] = None
    job_description: Optional[str] = None


# --- 2. AI AGENT STRATEGY MODELS ---

class ResumeStrategy(BaseModel):
    """Output schema for the Supervisor Agent's strategic planning phase."""
    reasoning_summary: str = Field(
        description="The internal logic explaining the pivot strategy."
    )
    gaps_identified: List[Union[str, dict]] = Field(
        description="List of detected skill deficiencies or narrative weaknesses."
    )
    instruction_to_writer: str = Field(
        description="Technical directives for the Writer Agent to follow."
    )
    next_action: str = Field(
        description="Determines the next node in the graph (e.g., 'delegate_to_writer')."
    )


# --- 3. ANALYTICS & SCORING MODELS ---

class JobMatchReport(BaseModel):
    """Schema for ATS-style alignment analysis against a specific job role."""
    score: int = Field(
        description="Numerical match assessment ranging from 0 to 100."
    )
    missing_keywords: List[str] = Field(
        description="Specific technical or soft skills missing from the resume."
    )
    hiring_manager_tip: str = Field(
        description="Actionable advice to make the candidate more competitive."
    )

"""
Thrive Wellness Career Transition Platform - Application Entry Point
------------------------------------------------------------------
This module serves as the frontend interface for the multi-agent career 
transition system. It handles:
1. User input and profile management via Streamlit sidebar.
2. Local job database loading and vector search indexing.
3. Orchestration of Agent Graphs (Resume, Application, Search).
4. Local LaTeX compilation and PDF rendering.
5. PII Restoration and data privacy enforcement.
"""

import streamlit as st
import subprocess
import os
import json
import base64
import re
import urllib.parse
import chromadb
from chromadb.utils import embedding_functions
from dotenv import load_dotenv

load_dotenv()
from schemas import CandidateProfile, Experience, Education, Project 
from agent_graph import build_graph, build_application_graph, build_job_search_graph

# --- PAGE CONFIGURATION ---
st.set_page_config(page_title="Thrive Career Wellness", page_icon="🌿", layout="wide")

# Initialize session state counters for dynamic UI elements
if 'exp_count' not in st.session_state: st.session_state.exp_count = 1
if 'proj_count' not in st.session_state: st.session_state.proj_count = 1

# --- DATA LOADING ---
@st.cache_data
def load_jobs():
    """
    Loads the job market database from local storage.

    Returns:
        list: A list of job dictionaries. Returns a default placeholder if 
              the database file is missing.
    """
    try:
        with open("jobs_database.json", "r") as f: 
            return json.load(f)
    except FileNotFoundError: 
        return [{"id": 0, "title": "Default", "sector": "Tech", "description": "N/A"}]

@st.cache_data
def load_candidates():
    """
    Loads pre-defined candidate personas for demonstration purposes.

    Returns:
        list: A list of candidate profile dictionaries.
    """
    try:
        with open("candidates_database.json", "r") as f: 
            return json.load(f)
    except FileNotFoundError: 
        return []

jobs_db = load_jobs()
candidates_db = load_candidates()

# --- VECTOR DATABASE ---
@st.cache_resource
class JobBoard:
    """
    Manages vector search operations for job market discovery.

    Attributes:
        client: The ChromaDB persistent client.
        collection: The specific document collection for job embeddings.
    """
    def __init__(self):
        self.client = chromadb.PersistentClient(path="./chroma_db")
        self.collection = self.client.get_or_create_collection(
            name="job_market", 
            embedding_function=embedding_functions.DefaultEmbeddingFunction()
        )
        if self.collection.count() == 0: 
            self._index_jobs()

    def _index_jobs(self):
        """
        Populates the vector database with job descriptions from the local store.
        """
        ids = [str(j['id']) for j in jobs_db]
        docs = [f"{j['title']} - {j['description']}" for j in jobs_db]
        metadatas = [{"title": j['title'], "sector": j.get('sector', 'General'), "id": j['id']} for j in jobs_db]
        self.collection.add(documents=docs, metadatas=metadatas, ids=ids)

    def recommend_jobs(self, resume_text, top_k=50):
        """
        Performs semantic similarity search against the job database.

        Args:
            resume_text (str): The candidate's resume content or summary.
            top_k (int): Number of results to retrieve.

        Returns:
            list: A list of job dictionaries matching the semantic query.
        """
        results = self.collection.query(query_texts=[resume_text], n_results=top_k)
        jobs = []
        if results['ids']:
            for i in range(len(results['ids'][0])):
                meta = results['metadatas'][0][i]
                jobs.append({
                    "id": meta['id'], 
                    "title": meta['title'], 
                    "sector": meta['sector'], 
                    "description": results['documents'][0][i]
                })
        return jobs

# --- LOCAL LATEX COMPILER ---
def compile_latex_local(latex_code):
    """
    Compiles LaTeX source into PDF using the local Tectonic engine.

    Args:
        latex_code (str): The raw LaTeX source string.

    Returns:
        tuple: (bytes, None) on success, or (None, str) containing stderr on failure.
    """
    with open("resume.tex", "w", encoding="utf-8") as f: 
        f.write(latex_code)
    try:
        res = subprocess.run(
            [os.path.join(os.getcwd(), "tectonic.exe"), "resume.tex"], 
            capture_output=True, 
            text=True            
        )
        if res.returncode == 0:
            with open("resume.pdf", "rb") as f: 
                return f.read(), None
        else:
            return None, res.stderr
    except Exception as e: 
        return None, str(e)

# --- SIDEBAR: PROFILE MANAGEMENT ---
with st.sidebar:
    st.header("1. Member Profile")
    persona_names = ["New Blank Profile"] + [f"{p.get('level', 'N/A')} - {p['full_name']}" for p in candidates_db]
    selected_persona = st.selectbox("Select Candidate:", persona_names)

    if selected_persona == "New Blank Profile":
        p_data = {"full_name": "", "contact_email": "", "phone": "", "linkedin": "", 
                  "experience_history": [], "projects": [], 
                  "education": [{"school": "", "degree": "", "year": ""}]}
        st.session_state.exp_count = 1
        st.session_state.proj_count = 1
    else:
        p_data = candidates_db[persona_names.index(selected_persona) - 1]
        st.session_state.exp_count = max(1, len(p_data.get('experience_history', [])))
        st.session_state.proj_count = max(1, len(p_data.get('projects', [])))

    full_name = st.text_input("Full Name", p_data['full_name'], key=f"fn_{selected_persona}")
    email = st.text_input("Email", p_data['contact_email'], key=f"em_{selected_persona}")
    phone = st.text_input("Phone", p_data['phone'], key=f"ph_{selected_persona}")
    linkedin = st.text_input("LinkedIn", p_data.get('linkedin', ""), key=f"li_{selected_persona}")

    with st.expander("💼 Work History", expanded=True):
        for i in range(st.session_state.exp_count):
            st.markdown(f"**Job #{i+1}**")
            d_exp = p_data['experience_history'][i] if i < len(p_data['experience_history']) else {"company": "", "role": "", "achievements": [""]}
            st.text_input("Company", key=f"comp_{selected_persona}_{i}", value=d_exp['company'])
            st.text_input("Role", key=f"role_{selected_persona}_{i}", value=d_exp['role'])
            st.text_area("Impact Story", key=f"ach_{selected_persona}_{i}", value="\n".join(d_exp['achievements']))
        if st.button("➕ Add Job"): 
            st.session_state.exp_count += 1
            st.rerun()

    with st.expander("🚀 Projects", expanded=False):
        for j in range(st.session_state.proj_count):
            st.markdown(f"**Project #{j+1}**")
            d_pj = p_data['projects'][j] if j < len(p_data['projects']) else {"title": "", "description_bullets": [""]}
            st.text_input("Title", key=f"ptit_{selected_persona}_{j}", value=d_pj['title'])
            st.text_area("Story", key=f"pdesc_{selected_persona}_{j}", value="\n".join(d_pj['description_bullets']))
        if st.button("➕ Add Project"): 
            st.session_state.proj_count += 1
            st.rerun()

    with st.expander("🎓 Education", expanded=False):
        d_edu = p_data['education'][0] if p_data.get('education') else {"school": "", "degree": "", "year": ""}
        school = st.text_input("University", d_edu['school'], key=f"sch_{selected_persona}")
        degree = st.text_input("Degree", d_edu['degree'], key=f"deg_{selected_persona}")
        grad_year = st.text_input("Year", d_edu['year'], key=f"yr_{selected_persona}")

    st.header("2. Target Role")
    sects = sorted(list(set(j['sector'] for j in jobs_db)))
    sel_sect = st.selectbox("Industry Sector", sects)
    filt_jobs = [j for j in jobs_db if j['sector'] == sel_sect]
    sel_job_t = st.selectbox("Role Template", [f"{j['title']} (ID: {j['id']})" for j in filt_jobs])
    sel_job_obj = next(j for j in filt_jobs if f"{j['title']} (ID: {j['id']})" == sel_job_t)
    target_role = st.text_input("Target Role Name", sel_job_obj['title'])
    job_desc = st.text_area("Job Description", sel_job_obj['description'], height=150)
    
    submit = st.button("🚀 Generate Strategic Resume", type="primary")

# --- EXECUTION LAYER ---
if submit:
    # Collect UI Data
    exps = []
    for i in range(st.session_state.exp_count):
        k_comp, k_role, k_ach = f"comp_{selected_persona}_{i}", f"role_{selected_persona}_{i}", f"ach_{selected_persona}_{i}"
        c = st.session_state[k_comp] if k_comp in st.session_state else p_data['experience_history'][i]['company']
        r = st.session_state[k_role] if k_role in st.session_state else p_data['experience_history'][i]['role']
        a = st.session_state[k_ach] if k_ach in st.session_state else "\n".join(p_data['experience_history'][i]['achievements'])
        exps.append(Experience(company=c, location="N/A", role=r, duration="N/A", achievements=[a]))

    projs = []
    for j in range(st.session_state.proj_count):
        k_tit, k_desc = f"ptit_{selected_persona}_{j}", f"pdesc_{selected_persona}_{j}"
        t = st.session_state[k_tit] if k_tit in st.session_state else p_data['projects'][j]['title']
        d = st.session_state[k_desc] if k_desc in st.session_state else "\n".join(p_data['projects'][j]['description_bullets'])
        projs.append(Project(title=t, tech_stack="N/A", date="N/A", description_bullets=[d]))

    k_sch, k_deg, k_yr = f"sch_{selected_persona}", f"deg_{selected_persona}", f"yr_{selected_persona}"
    edu_school = st.session_state[k_sch] if k_sch in st.session_state else p_data['education'][0]['school']
    edu_degree = st.session_state[k_deg] if k_deg in st.session_state else p_data['education'][0]['degree']
    edu_year = st.session_state[k_yr] if k_yr in st.session_state else p_data['education'][0]['year']

    profile = CandidateProfile(
        full_name=full_name, contact_email=email, phone=phone, linkedin=linkedin, 
        summary=f"Strategic {target_role} transition.", skills=[], 
        experience_history=exps, projects=projs, 
        education=[Education(school=edu_school, location="N/A", degree=edu_degree, year=edu_year)], 
        target_role=target_role, job_description=job_desc
    )
    st.session_state.profile = profile

    # Orchestrate Multi-Agent Workflow
    with st.status("🤖 AI Agents Initializing...", expanded=True) as status:
        st.write("🧠 Reasoning through career pivot strategy...")
        app = build_graph()
        st.session_state.result = app.invoke({"profile": profile, "target_role": target_role})
        
        st.write("🎯 Mapping broad market opportunities...")
        board = JobBoard()
        raw = board.recommend_jobs(target_role, top_k=50)
        search_agent = build_job_search_graph()
        matches = search_agent.invoke({"profile": profile, "raw_jobs": raw, "ranked_jobs": []})
        st.session_state.ranked_results = matches['ranked_jobs']
        
        status.update(label="✅ Ready!", state="complete")
    st.rerun()

# --- PRESENTATION LAYER ---
if 'result' in st.session_state and st.session_state.result:
    res = st.session_state.result
    prof = st.session_state.profile
    
    st.divider()
    st.header("📋 Strategic Career Assets")
    
    with st.expander("🧠 View AI Reasoning & Pivot Strategy", expanded=True):
        st.write(res['strategy'].reasoning_summary)

    # PII RESTORATION: Double-pass loop to handle AI-escaped placeholders
    raw_latex = res['final_draft']
    pii_map = {
        "NAME": prof.full_name,
        "EMAIL": prof.contact_email,
        "NUMBER": prof.phone,
        "URL": prof.linkedin
    }

    clean_latex = raw_latex
    for key, value in pii_map.items():
        # Pass 1: Standard Placeholders
        clean_latex = clean_latex.replace(f"[CANDIDATE_{key}]", value)\
                                 .replace(f"[CONTACT_{key}]", value)\
                                 .replace(f"[PHONE_{key}]", value)\
                                 .replace(f"[LINKEDIN_{key}]", value)
        
        # Pass 2: Escaped Placeholders (Latex safe)
        clean_latex = clean_latex.replace(f"[CANDIDATE\\_{key}]", value)\
                                 .replace(f"[CONTACT\\_{key}]", value)\
                                 .replace(f"[PHONE\\_{key}]", value)\
                                 .replace(f"[LINKEDIN\\_{key}]", value)

    c1, c2 = st.columns(2)
    with c1:
        st.subheader("LaTeX Source Editor")
        edited = st.text_area("Source Code", value=clean_latex, height=600, key="latex_editor")
    with c2:
        st.subheader("Resume Preview")
        if st.button("📄 Compile PDF Document"):
            with st.spinner("Generating PDF..."):
                pdf_bytes, error_log = compile_latex_local(edited)
                
                if pdf_bytes:
                    b64 = base64.b64encode(pdf_bytes).decode()
                    st.markdown(f'<iframe src="data:application/pdf;base64,{b64}" width="100%" height="600"></iframe>', unsafe_allow_html=True)
                else:
                    st.error("LaTeX Compilation Failed")
                    with st.expander("🔍 View Compiler Error Log"):
                        st.code(error_log)

    if 'ranked_results' in st.session_state:
        st.divider()
        st.header("🎯 Thrive Job Match (Top 10)")
        m1, m2 = st.columns(2)
        for i, job in enumerate(st.session_state.ranked_results):
            job_id = job['id']
            with (m1 if i % 2 == 0 else m2):
                with st.container(border=True):
                    st.subheader(job['title'])
                    st.caption(f"Sector: {job['sector']} | ID: {job_id}")
                    with st.expander("🔍 Full Responsibilities & Details"): 
                        st.write(job['description'])
                    
                    ce, cp = st.columns([1, 1.2])
                    
                    with ce:
                        cl_base = st.session_state.get(f"packet_{job_id}", f"Hi, I'm {prof.full_name}, I'm writing to express interest in the {job['title']} role...")
                        cl_clean = cl_base.replace("[CANDIDATE_NAME]", prof.full_name)\
                                          .replace("[CONTACT_EMAIL]", prof.contact_email)
                        
                        subj = urllib.parse.quote(f"Interest in {job['title']}")
                        mailto = f"mailto:hiring@thrive.com?subject={subj}&body={urllib.parse.quote(cl_clean)}"
                        st.link_button("📧 Email Manager", mailto, use_container_width=True)
                    
                    with cp:
                        if st.button(f"🤖 Prep Cover Letter", key=f"prep_{job_id}", use_container_width=True):
                            with st.spinner("Tailoring cover letter..."):
                                
                                # 🛡️ PII GUARD: Create anonymized profile for the agent
                                safe_prof = prof.model_copy()
                                safe_prof.full_name = "[CANDIDATE_NAME]"
                                safe_prof.contact_email = "[CONTACT_EMAIL]"
                                safe_prof.phone = "[PHONE_NUMBER]"
                                safe_prof.linkedin = "[LINKEDIN_URL]"

                                app_res = build_application_graph().invoke({
                                    "job_title": job['title'], 
                                    "job_description": job['description'], 
                                    "candidate_profile": safe_prof.model_dump_json(), # Safe JSON
                                    "cover_letter": "", "status": "pending"
                                })
                                st.session_state[f"packet_{job_id}"] = app_res['cover_letter']
                                st.rerun()

                    if f"packet_{job_id}" in st.session_state:
                        st.info("📑 Tailored Application Ready")
                        final_cl = st.session_state[f"packet_{job_id}"]
                        for key, value in pii_map.items():
                             final_cl = final_cl.replace(f"[CANDIDATE_{key}]", value)\
                                                .replace(f"[CONTACT_{key}]", value)
                        st.text_area("Cover Letter", value=final_cl, height=200, key=f"text_{job_id}")