import streamlit as st
from dotenv import load_dotenv
from pathlib import Path
from typing import Any, Dict, List, Optional
from llama_index.llama_pack.base import BaseLlamaPack
from llama_index.readers import PDFReader
from llama_index.llms.base import LLM
from llama_index.llms import OpenAI
from llama_index import ServiceContext
from llama_index.schema import NodeWithScore
from llama_index.response_synthesizers import TreeSummarize
from pydantic import BaseModel
import os
import pdfplumber
import io

# Load environment variables from .env file
load_dotenv()

# Get OpenAI API key from environment variables
openai_api_key = os.getenv("OPENAI_API_KEY")

QUERY_TEMPLATE = """
You are an expert resume reviewer. 
Your job is to decide if the candidate passes the resume screen given the job description and a list of criteria:

### Job Description
{job_description}

### Screening Criteria
{criteria_str}
"""

class CriteriaDecision(BaseModel):
    """The decision made based on a single criterion"""
    decision: bool
    reasoning: str

class ResumeScreenerDecision(BaseModel):
    """The decision made by the resume screener"""
    criteria_decisions: List[CriteriaDecision]
    overall_reasoning: str
    overall_decision: bool

def _format_criteria_str(criteria: List[str]) -> str:
    criteria_str = ""
    for criterion in criteria:
        criteria_str += f"- {criterion}\n"
    return criteria_str

class ResumeScreenerPack(BaseLlamaPack):
    def _init_(
        self, job_description: str, criteria: List[str], llm: Optional[LLM] = None
    ) -> None:
        self.reader = PDFReader()
        llm = llm or OpenAI(model="gpt-4", api_key=openai_api_key)
        service_context = ServiceContext.from_defaults(llm=llm)
        self.synthesizer = TreeSummarize(
            output_cls=ResumeScreenerDecision, service_context=service_context
        )
        criteria_str = _format_criteria_str(criteria)
        self.query = QUERY_TEMPLATE.format(
            job_description=job_description, criteria_str=criteria_str
        )

    def get_modules(self) -> Dict[str, Any]:
        """Get modules."""
        return {"reader": self.reader, "synthesizer": self.synthesizer}

    def run(self, resume_path: str, *args: Any, **kwargs: Any) -> Any:
        """Run pack."""
        docs = self.reader.load_data(Path(resume_path))
        output = self.synthesizer.synthesize(
            query=self.query,
            nodes=[NodeWithScore(node=doc, score=1.0) for doc in docs],
        )
        return output.response

def main():
    st.title("Resume Screener App")

    # Sidebar for user input
    job_description = st.text_area("Job Description")
    criteria = st.text_area("Screening Criteria (separate each criterion by a new line)")

    uploaded_file = st.file_uploader("Upload Resume (PDF)", type=["pdf"])

    if st.button("Submit"):
        if job_description and criteria and uploaded_file:
            resume_text = extract_text_from_pdf(uploaded_file)

            screener_pack = ResumeScreenerPack(job_description=job_description, criteria=criteria.split("\n"))

            with st.spinner("Analyzing the resume..."):
                result = screener_pack.run(resume_text)

            st.subheader("Screening Results")
            st.json(result)

def extract_text_from_pdf(uploaded_file):
    if uploaded_file is not None:
        try:
            # Read PDF content from BytesIO
            uploaded_content = io.BytesIO(uploaded_file.read())

            with pdfplumber.open(uploaded_content) as pdf:
                text = ""
                for page in pdf.pages:
                    text += page.extract_text()
                return text
        except Exception as e:
            st.error(f"Error extracting text from PDF: {str(e)}")
            return ""
    else:
        st.error("Please upload a PDF file.")
        return ""

if _name_ == "_main_":
    main()