from fpdf import FPDF
from datetime import datetime

class PDFReport(FPDF):
    def header(self):
        self.set_font('Arial', 'B', 14)
        self.cell(0, 10, 'Lung Disease Detection Report', ln=True, align='C')
        self.ln(10)

    def add_patient_info(self, name, age, gender):
        self.set_font('Arial', '', 12)
        self.cell(0, 10, f'Name: {name}    Age: {age}    Gender: {gender}', ln=True)
        self.cell(0, 10, f'Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}', ln=True)
        self.ln(5)

    def add_images(self, xray_path, mask_path):
        self.cell(0, 10, 'X-ray Image:', ln=True)
        self.image(xray_path, w=90)
        self.ln(5)
        self.cell(0, 10, 'Segmented Output:', ln=True)
        self.image(mask_path, w=90)
        self.ln(10)

    def add_diagnosis(self, disease, severity, comment):
        self.set_font('Arial', 'B', 12)
        self.cell(0, 10, f'Diagnosis: {disease}', ln=True)
        self.cell(0, 10, f'Severity: {severity}', ln=True)
        self.set_font('Arial', '', 12)
        self.multi_cell(0, 10, f'Comments:\n{comment}')
        self.ln(10)

    def footer(self):
        self.set_y(-15)
        self.set_font('Arial', 'I', 10)
        self.cell(0, 10, 'AI-generated report. Please consult a certified radiologist.', align='C')


def get_disease_comment(disease, severity):
    comments = {
        "COVID-19": {
            "Mild": "Ground-glass opacities are faintly visible, typically associated with mild COVID-19. Recommend monitoring and follow-up scan.",
            "Moderate": "Bilateral opacities with peripheral distribution indicate moderate COVID-19 pneumonia. Suggest home isolation and treatment unless symptoms worsen.",
            "Severe": "Extensive ground-glass opacities and consolidation detected. Suggest immediate hospitalization for advanced COVID-19 treatment."
        },
        "Pneumonia": {
            "Mild": "Localized patchy infiltrates suggest early-stage pneumonia. Recommend antibiotics and follow-up.",
            "Moderate": "Lung lobes show dense consolidation, typical of moderate pneumonia. Clinical treatment and rest advised.",
            "Severe": "Widespread opacities indicate severe pneumonia. Urgent clinical attention required."
        },
        "Tuberculosis": {
            "Mild": "Apical scarring may indicate early-stage tuberculosis. Recommend sputum test and anti-TB medication.",
            "Moderate": "Fibronodular lesions seen, likely due to active TB. Suggest starting anti-tubercular therapy immediately.",
            "Severe": "Extensive cavitary lesions observed, typical of severe TB. Hospitalization and long-term treatment needed."
        },
        "Normal": {
            "Mild": "Lungs appear healthy. No abnormal findings.",
            "Moderate": "Slight deviations detected but not indicative of major disease.",
            "Severe": "Abnormal scan reported, but not matching major disease pattern. Recommend detailed physical exam."
        }
    }

    return comments.get(disease, {}).get(severity, "No specific comment available.")


def generate_report(name, age, gender, xray_path, mask_path, disease, severity, output_path):
    comment = get_disease_comment(disease, severity)
    pdf = PDFReport()
    pdf.add_page()
    pdf.add_patient_info(name, age, gender)
    pdf.add_images(xray_path, mask_path)
    pdf.add_diagnosis(disease, severity, comment)
    pdf.output(output_path)