import os
import docx
from docxtpl import DocxTemplate
from typing import TypedDict, List
from pydantic import BaseModel, Field
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.graph import StateGraph, START, END
import logging

logger = logging.getLogger(__name__)

# --- Define State & Schemas ---
class ResumeState(TypedDict):
    original_resume: str
    job_description: str
    resume_skills: List[str]
    jd_skills: List[str]
    matched_skills: List[str]
    tailored_resume: dict

class SkillList(BaseModel):
    skills: List[str] = Field(description="A list of technical skills, languages, frameworks, and databases")

class ExperienceItem(BaseModel):
    job_title: str = Field(description="The job title")
    company: str = Field(description="The company name")
    duration: str = Field(description="The time period worked")
    bullets: List[str] = Field(description="3-5 bullet points emphasizing the matched skills")

class EducationItem(BaseModel):
    degree: str = Field(description="The name of the degree or certification")
    institution: str = Field(description="The name of the university or institution")
    graduation_date: str = Field(description="The graduation year or timeframe")

class ATSResumeOutput(BaseModel):
    name: str = Field(description="The candidate's full name")
    contact_info: str = Field(description="Email, phone, and links (LinkedIn, GitHub)")
    summary: str = Field(description="A professional summary emphasizing matched skills")
    skills: List[str] = Field(description="The optimized list of relevant skills")
    experience: List[ExperienceItem] = Field(description="The candidate's work history")
    education: List[EducationItem] = Field(description="The candidate's educational background")

class AILangGraphTailor:
    def __init__(self):
        # Initialize Gemini 2.5 Flash
        # Assumes GOOGLE_API_KEY is loaded in the environment 
        self.llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0.2)
        self.app = self._build_graph()

    def extract_text_from_docx(self, file_path: str) -> str:
        """Reads a .docx file and returns all text as a single string."""
        try:
            doc = docx.Document(file_path)
            full_text = [para.text for para in doc.paragraphs if para.text.strip()]
            return "\n".join(full_text)
        except Exception as e:
            logger.error(f"Error reading document: {e}")
            raise

    def generate_final_docx(self, tailored_dict: dict, template_path: str, output_path: str):
        """Takes the JSON dictionary from Gemini and renders the final Word document using docxtpl."""
        try:
            doc = DocxTemplate(template_path)
            doc.render(tailored_dict)
            doc.save(output_path)
            logger.info(f"Success! Tailored document saved to {output_path}.")
        except Exception as e:
            logger.error(f"Error rendering docxtpl: {e}")
            raise

    # --- Node Functions ---
    def extract_resume_skills(self, state: ResumeState):
        structured_llm = self.llm.with_structured_output(SkillList)
        prompt = f"Extract all technical skills from this resume. Return only the skills.\n\nResume:\n{state['original_resume']}"
        return {"resume_skills": structured_llm.invoke(prompt).skills}

    def extract_jd_skills(self, state: ResumeState):
        structured_llm = self.llm.with_structured_output(SkillList)
        prompt = f"Extract all required technical skills from this job description. Return only the skills.\n\nJob Description:\n{state['job_description']}"
        return {"jd_skills": structured_llm.invoke(prompt).skills}

    def match_skills(self, state: ResumeState):
        resume_set = {s.strip().lower() for s in state['resume_skills']}
        jd_set = {s.strip().lower() for s in state['jd_skills']}
        matches = list(resume_set.intersection(jd_set))
        return {"matched_skills": matches}

    def tailor_resume(self, state: ResumeState):
        matched_skills_str = ", ".join(state['matched_skills'])
        structured_llm = self.llm.with_structured_output(ATSResumeOutput)
        prompt = f"""
        You are a professional resume writer. Your task is to rewrite the provided 'Original Resume' to strongly emphasize these 'Key Skills to Emphasize': {matched_skills_str}

        RULES:
        1. Rephrase the summary and experience bullet points to highlight the key skills, and do not use buzzwords while writing any Summary and experience. Keep everything simple.
        2. You MUST NOT add any new skills, projects, experiences, or education that are not in the Original Resume.
        3. Ensure the output is a professional, complete resume.

        Original Resume:
        ---
        {state['original_resume']}
        ---
        """
        tailored_data = structured_llm.invoke(prompt)
        return {"tailored_resume": tailored_data.model_dump()}

    def _build_graph(self):
        workflow = StateGraph(ResumeState)
        workflow.add_node("extract_resume", self.extract_resume_skills)
        workflow.add_node("extract_jd", self.extract_jd_skills)
        workflow.add_node("match_skills", self.match_skills)
        workflow.add_node("tailor_resume", self.tailor_resume)

        workflow.add_edge(START, "extract_resume")
        workflow.add_edge("extract_resume", "extract_jd")
        workflow.add_edge("extract_jd", "match_skills")
        workflow.add_edge("match_skills", "tailor_resume")
        workflow.add_edge("tailor_resume", END)
        return workflow.compile()

    def run_pipeline(self, original_resume_text: str, job_description: str) -> dict:
        """Executes the LangGraph pipeline."""
        initial_state = {
            "original_resume": original_resume_text,
            "job_description": job_description
        }
        return self.app.invoke(initial_state)