import os
import time
import json
import requests
from typing import List, Optional, Dict, Any
from pydantic import BaseModel, Field
from mistralai import Mistral

import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

os.environ['MISTRAL_API_KEY'] = "YOUR MISTRALAI API KEY" # Get it from https://console.mistral.ai/api-keys

client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])

def download_job_description(url, output_path = "job_description.pdf"):
    """
    Download job description from a given URL.
    """
    response = requests.get(url)
    with open(output_path, "wb") as f:
        f.write(response.content)
    print(f"Downloaded {output_path}")

def download_resumes(url, local_dir="resumes"):
    """
    Download resumes from the given URL.
    """

    response = requests.get(url)
    
    if response.status_code != 200:
        print("Failed to retrieve folder contents:", response.text)
        return
    
    data = response.json()
    os.makedirs(local_dir, exist_ok=True)

    print(f"{len(data)} files available for download:")
    for file in data:
        file_name = file["name"]
        download_url = file["download_url"]

        r = requests.get(download_url)
        with open(os.path.join(local_dir, file_name), "wb") as f:
            f.write(r.content)
        print(f"Downloaded {file_name}")

class Skill(BaseModel):
    name: str = Field(description="Name of the skill or technology")
    level: Optional[str] = Field(description="Proficiency level (beginner, intermediate, advanced)")
    years: Optional[float] = Field(description="Years of experience with this skill")

class Education(BaseModel):
    degree: str = Field(description="Type of degree or certification obtained")
    field: str = Field(description="Field of study or specialization")
    institution: str = Field(description="Name of educational institution")
    year_completed: Optional[int] = Field(description="Year when degree was completed")
    gpa: Optional[float] = Field(description="Grade Point Average, typically on 4.0 scale")

class Experience(BaseModel):
    title: str = Field(description="Job title or position held")
    company: str = Field(description="Name of employer or organization")
    duration_years: float = Field(description="Duration of employment in years")
    skills_used: List[str] = Field(description="Skills utilized in this role")
    achievements: List[str] = Field(description="Key accomplishments or responsibilities")
    relevance_score: Optional[float] = Field(description="Relevance to current job opening (0-10 scale)")

class ContactDetails(BaseModel):
    name: str = Field(description="Full name of the candidate")
    email: str = Field(description="Primary email address for contact")
    phone: Optional[str] = Field(description="Phone number with country code if applicable")
    location: Optional[str] = Field(description="Current city and country/state")
    linkedin: Optional[str] = Field(description="LinkedIn profile URL")
    website: Optional[str] = Field(description="Personal or portfolio website URL")

class JobRequirements(BaseModel):
    required_skills: List[Skill] = Field(description="Skills that are mandatory for the position")
    preferred_skills: List[Skill] = Field(description="Skills that are desired but not required")
    min_experience_years: float = Field(description="Minimum years of experience required")
    required_education: List[Education] = Field(description="Mandatory educational qualifications")
    preferred_domains: List[str] = Field(description="Industry domains or sectors preferred for experience")

class CandidateProfile(BaseModel):
    contact_details: ContactDetails = Field(description="Candidate's personal and contact information")
    skills: List[Skill] = Field(description="Technical and soft skills possessed by the candidate")
    education: List[Education] = Field(description="Educational background and qualifications")
    experience: List[Experience] = Field(description="Professional work history and experience")

class SkillMatch(BaseModel):
    skill_name: str = Field(description="Name of the skill being evaluated")
    present: bool = Field(description="Whether the candidate possesses this skill")
    match_level: float = Field(description="How well the candidate's skill matches the requirement (0-10 scale)")
    confidence: float = Field(description="Confidence in the skill evaluation (0-1 scale)")
    notes: str = Field(description="Additional context about the skill match assessment")

class CandidateScore(BaseModel):
    technical_skills_score: float = Field(description="Assessment of technical capabilities (0-40 points)")
    experience_score: float = Field(description="Evaluation of relevant work experience (0-30 points)")
    education_score: float = Field(description="Rating of educational qualifications (0-15 points)")
    additional_score: float = Field(description="Score for other relevant factors (0-15 points)")
    total_score: float = Field(description="Aggregate candidate evaluation score (0-100)")
    key_strengths: List[str] = Field(description="Primary candidate advantages for this role")
    key_gaps: List[str] = Field(description="Areas where the candidate lacks desired qualifications")
    confidence: float = Field(description="Overall confidence in the evaluation accuracy (0-1 scale)")
    notes: str = Field(description="Supplementary observations about the candidate fit")

class CandidateResult(BaseModel):
    file_name: str = Field(description="Name of the source resume file")
    contact_details: ContactDetails = Field(description="Candidate's contact information")
    candidate_profile: CandidateProfile = Field(description="Complete extracted candidate profile")
    score: CandidateScore = Field(description="Detailed evaluation scores and assessment")

class Agent:
    def __init__(self, name: str, client: Mistral):
        self.name = name
        self.client = client

    def process(self, message):
        """Base process method - to be implemented by child classes"""
        raise NotImplementedError("Subclasses must implement process method")

    def communicate(self, recipient_agent, message):
        """Send message to another agent"""
        return recipient_agent.process(message)

class DocumentAgent(Agent):
    def __init__(self, client: Mistral):
        super().__init__("DocumentAgent", client)

    def process(self, file_info):
        """Process document extraction request"""
        file_path, file_name = file_info
        return self.extract_text_from_file(file_path, file_name)

    def extract_text_from_file(self, file_path: str, file_name: str) -> str:
        """Extract text from a file using Mistral OCR"""
        try:
            # Upload the file
            uploaded_file = self.client.files.upload(
                file={
                    "file_name": file_name,
                    "content": open(file_path, "rb"),
                },
                purpose="ocr"
            )

            # Get signed URL
            signed_url = self.client.files.get_signed_url(file_id=uploaded_file.id)

            # Process with OCR
            ocr_response = self.client.ocr.process(
                model="mistral-ocr-latest",
                document={
                    "type": "document_url",
                    "document_url": signed_url.url,
                }
            )

            # Extract and return the text
            extracted_text = ""
            for page in ocr_response.pages:
                extracted_text += page.markdown + "\n\n"

            return extracted_text

        except Exception as e:
            print(f"Error extracting text from {file_name}: {str(e)}")
            return ""

class JobAnalysisAgent(Agent):
    def __init__(self, client: Mistral):
        super().__init__("JobAnalysisAgent", client)

    def process(self, jd_text):
        """Process job description text"""
        return self.extract_job_requirements(jd_text)

    def extract_job_requirements(self, jd_text: str) -> JobRequirements:
        """Extract structured job requirements from a job description"""
        prompt = f"""
        Extract the key job requirements from the following job description.
        Focus on required skills, preferred skills, experience requirements, and education requirements.

        Job Description:
        {jd_text}
        """

        response = self.client.chat.parse(
            model="mistral-small-latest",
            messages=[
                {"role": "system", "content": "Extract structured job requirements from the job description."},
                {"role": "user", "content": prompt}
            ],
            response_format=JobRequirements,
            temperature=0
        )

        return json.loads(response.choices[0].message.content)

class ResumeAnalysisAgent(Agent):
    def __init__(self, client: Mistral):
        super().__init__("ResumeAnalysisAgent", client)

    def process(self, resume_text):
        """Process resume text"""
        return self.extract_candidate_profile(resume_text)

    def extract_candidate_profile(self, resume_text: str) -> CandidateProfile:
        """Extract structured candidate information from resume text"""
        prompt = f"""
        Extract the candidate's contact details, skills, education, and experience from the following resume.
        Be thorough and include all relevant information.

        Resume:
        {resume_text}
        """

        response = self.client.chat.parse(
            model="mistral-small-latest",
            messages=[
                {"role": "system", "content": "Extract structured candidate information from the resume."},
                {"role": "user", "content": prompt}
            ],
            response_format=CandidateProfile,
            temperature=0
        )

        return json.loads(response.choices[0].message.content)

class MatchingAgent(Agent):
    def __init__(self, client: Mistral):
        super().__init__("MatchingAgent", client)

    def process(self, data):
        """Process job requirements and candidate profile to generate score"""
        job_requirements, candidate_profile, resume_text = data
        return self.evaluate_candidate(job_requirements, candidate_profile, resume_text)

    def evaluate_candidate(self, job_requirements: JobRequirements, candidate_profile: CandidateProfile, resume_text: str) -> CandidateScore:
        """Evaluate how well a candidate matches the job requirements"""
        # Convert to JSON for inclusion in the prompt
        job_req_json = json.dumps(job_requirements, indent=2)
        candidate_json = json.dumps(candidate_profile, indent=2)

        prompt = f"""
        Evaluate how well the candidate matches the job requirements.

        Job Requirements:
        {job_req_json}

        Candidate Profile:
        {candidate_json}

        Provide a detailed scoring breakdown, highlighting strengths and gaps.
        Assess the quality and relevance of the candidate's experience, not just keyword matches.
        Include confidence levels for your assessment.

        Technical skills should be scored out of 40 points.
        Experience should be scored out of 30 points.
        Education should be scored out of 15 points.
        Additional qualifications should be scored out of 15 points.
        The total score should be out of 100 points.
        """

        response = self.client.chat.parse(
            model="mistral-small-latest",
            messages=[
                {"role": "system", "content": "Evaluate the candidate's match to the job requirements with detailed scoring."},
                {"role": "user", "content": prompt}
            ],
            response_format=CandidateScore,
            temperature=0.2  # Slight randomness for nuanced evaluation
        )

        return json.loads(response.choices[0].message.content)

class EmailCommunicationAgent(Agent):
    def __init__(self, client: Mistral, sender_email: str, app_password: str):
        super().__init__("EmailCommunicationAgent", client)
        self.sender_email = sender_email
        self.app_password = app_password

    def process(self, data):
        """Process email sending request"""
        candidate, calendly_link, subject = data
        return self.send_interview_invitation(candidate, calendly_link, subject)

    def send_interview_invitation(self, candidate, calendly_link: str, subject: str):
        """Generate and send personalized email to candidate"""
        name = candidate["contact_details"]['name']
        email = candidate["contact_details"]['email']

        # Create email HTML content
        html_content = f"""\
        <html>
          <body>
            <p>Hello {name},</p>
            <p>I'm the Hiring Manager from HireFive. Thank you for applying for the Data Scientist position at our company.</p>
            <p>We were impressed with your background and would like to schedule an initial screening call to discuss your experience and interest in the role.</p>
            <p>Please select a suitable time slot using our <a href="{calendly_link}">Calendly link</a>.</p>
            <p>Looking forward to speaking with you soon.</p>
            <p>Best regards,<br>
            Hiring Manager<br>
            HireFive</p>
          </body>
        </html>
        """

        if self.app_password:
            try:
                self.send_email(email, subject, html_content)
                return f"Email sent to {name} at {email}"
            except Exception as e:
                return f"Failed to send email to {name} ({email}): {str(e)}"
        else:
            return f"Would send email to {name} at {email} - Email subject: {subject}"

    def send_email(self, receiver_email, subject, html_content):
        """Send an email using Gmail SMTP"""
        # Create message container
        message = MIMEMultipart('alternative')
        message['From'] = self.sender_email
        message['To'] = receiver_email
        message['Subject'] = subject

        # Attach HTML part
        message.attach(MIMEText(html_content, 'html'))

        try:
            # Create SMTP session
            server = smtplib.SMTP('smtp.gmail.com', 587)
            server.starttls()  # Enable security

            # Login with Gmail account and app password
            server.login(self.sender_email, self.app_password)

            # Send email
            text = message.as_string()
            server.sendmail(self.sender_email, receiver_email, text)

        finally:
            server.quit()  # Close the connection

class CoordinatorAgent(Agent):
    def __init__(self, client: Mistral):
        super().__init__("CoordinatorAgent", client)
        self.document_agent = DocumentAgent(client)
        self.job_analysis_agent = JobAnalysisAgent(client)
        self.resume_analysis_agent = ResumeAnalysisAgent(client)
        self.matching_agent = MatchingAgent(client)
        self.email_communication_agent = None  # Will be initialized later with email credentials

    def set_email_communication_agent(self, sender_email: str, app_password: str):
        """Initialize communication agent with email credentials"""
        self.email_communication_agent = EmailCommunicationAgent(self.client, sender_email, app_password)

    def process_hiring_workflow(self, jd_file_path: str, resume_dir: str, output_path: str,
                               threshold_score: float, calendly_link: str, email_subject: str):
        """
        Coordinate the entire hiring workflow from document processing to interview scheduling
        """
        results = []

        # Process job description
        print(f"🤖 DocumentAgent extracting text from job description...")
        jd_text = self.document_agent.process((jd_file_path, os.path.basename(jd_file_path)))

        if not jd_text:
            print("❌ Failed to extract text from job description. Aborting.")
            return results

        # Extract job requirements
        print(f"🤖 JobAnalysisAgent analyzing job description...")
        job_requirements = self.job_analysis_agent.process(jd_text)

        time.sleep(10)

        # Process each resume in the directory
        resume_files = [f for f in os.listdir(resume_dir) if os.path.isfile(os.path.join(resume_dir, f))]

        for filename in resume_files[:5]:
            file_path = os.path.join(resume_dir, filename)
            print(f"\n🤖 DocumentAgent processing resume: {filename}")

            # Extract text from resume
            resume_text = self.document_agent.process((file_path, filename))

            time.sleep(10)

            if resume_text:
                # Extract candidate profile
                print(f"🤖 ResumeAnalysisAgent extracting candidate profile...")
                candidate_profile = self.resume_analysis_agent.process(resume_text)

                # Evaluate candidate match
                print(f"🤖 MatchingAgent evaluating candidate {candidate_profile['contact_details']['name']}...")
                score = self.matching_agent.process((job_requirements, candidate_profile, resume_text))

                # Create result object
                result = {
                    "file_name": filename,
                    "contact_details": candidate_profile["contact_details"],
                    "candidate_profile": candidate_profile,
                    "score": score
                }

                results.append(result)

                # Add a small delay to avoid rate limits
                time.sleep(10)
            else:
                print(f"❌ DocumentAgent failed to extract text from {filename}. Skipping this resume.")

        # Sort results by total score
        results.sort(key=lambda x: x["score"]['total_score'], reverse=True)

        # Save results to file
        with open(output_path, 'w') as f:
            json.dump([result for result in results], f, indent=2)

        print(f"\n🤖 CoordinatorAgent saved results to {output_path}")

        # Print summary of results
        print("\n===== CANDIDATE RANKING =====")
        for i, result in enumerate(results, 1):
            name = result["contact_details"]['name']
            score = result["score"]['total_score']
            print(f"{i}. {name}: {score}/100")

        # Send interview invitations to candidates above threshold
        if self.email_communication_agent:
            selected_candidates = [r for r in results if r["score"]['total_score'] >= threshold_score]

            print(f"\n🤖 EmailCommunicationAgent preparing to send interview invitations to {len(selected_candidates)} candidates who scored {threshold_score}+ out of 100...\n")

            for candidate in selected_candidates:
                response = self.email_communication_agent.process((candidate, calendly_link, email_subject))
                time.sleep(1)

        return results

jd_file_path = "job_description.pdf"
resume_dir = "resumes/"
output_path = "candidate_results.json"

sender_email = "<Your EmailID>"
app_password = "<Your generated app password>"
calendly_link = "<Your Calendly Link>"
email_subject = "HireFive: Next Steps for Your Data Scientist Application"

coordinator = CoordinatorAgent(client)

coordinator.set_communication_agent(sender_email, app_password)