import os
import json
import time
import re
import uuid
from datetime import datetime

from groq import Groq
from transformers import pipeline
from bert_score import score as bert_score
from rouge_score import rouge_scorer
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, HRFlowable
from reportlab.lib.enums import TA_CENTER
import gradio as gr


# ── Groq Client ──────────────────────────────────────────────
class GroqClientManager:
    def __init__(self):
        self.keys = []
        self.current_index = 0
        self._load_keys()

    def _load_keys(self):
        for i in range(1, 6):
            key = os.environ.get(f"GROQ_KEY_{i}")
            if key:
                self.keys.append(key)
        if not self.keys:
            raise ValueError("No Groq API keys found. Add GROQ_KEY_1 to Space secrets.")
        print(f"Loaded {len(self.keys)} Groq key(s)")

    def get_client(self):
        return Groq(api_key=self.keys[self.current_index])

    def rotate(self):
        self.current_index = (self.current_index + 1) % len(self.keys)

    def chat(self, messages, max_tokens=1500, temperature=0.3):
        for _ in range(len(self.keys)):
            try:
                client = self.get_client()
                response = client.chat.completions.create(
                    model="llama-3.3-70b-versatile",
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=temperature
                )
                return response.choices[0].message.content
            except Exception as e:
                if "429" in str(e) or "rate_limit" in str(e).lower():
                    self.rotate()
                    time.sleep(2)
                else:
                    raise e
        raise RuntimeError("All Groq API keys exhausted.")


groq_manager = GroqClientManager()

print("Loading NLI model...")
nli_pipeline = pipeline(
    "text-classification",
    model="cross-encoder/nli-deberta-v3-base"
)
print("NLI model loaded.")


# ── Report Prompts ────────────────────────────────────────────
REPORT_PROMPTS = {
    "radiology": """You are a radiologist writing a formal radiology report.
Given the following patient data, generate a structured radiology report.

Patient Data:
{patient_data}

Generate a report with these exact sections:
CLINICAL INDICATION:
TECHNIQUE:
FINDINGS:
IMPRESSION:

Be specific, clinical, and only use information provided. Do not invent findings.""",

    "discharge": """You are a hospital physician writing a discharge summary.
Given the following patient data, generate a structured discharge summary.

Patient Data:
{patient_data}

Generate a report with these exact sections:
PATIENT INFORMATION:
ADMISSION DIAGNOSIS:
HOSPITAL COURSE:
DISCHARGE DIAGNOSIS:
DISCHARGE MEDICATIONS:
FOLLOW-UP INSTRUCTIONS:

Only use information provided. Do not invent medications or diagnoses.""",

    "lab": """You are a clinical pathologist writing a laboratory report.
Given the following patient data, generate a structured lab report.

Patient Data:
{patient_data}

Generate a report with these exact sections:
TEST ORDERED:
SPECIMEN:
RESULTS:
REFERENCE RANGES:
INTERPRETATION:
RECOMMENDATION:

Only use information provided. Do not invent lab values."""
}


# ── Core Functions ────────────────────────────────────────────
def generate_report(patient_data, report_type):
    patient_data_str = "\n".join([f"{k}: {v}" for k, v in patient_data.items()])
    prompt = REPORT_PROMPTS[report_type].format(patient_data=patient_data_str)
    response = groq_manager.chat([{"role": "user", "content": prompt}])
    return {
        "report_type": report_type,
        "report_text": response.strip(),
        "patient_data": patient_data,
        "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }


def extract_sentences(text):
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s.strip() for s in sentences
            if len(s.strip()) > 20 and not s.strip().isupper() and ":" not in s[:25]]


def check_hallucination(report):
    source_text = " ".join([f"{k} is {v}." for k, v in report["patient_data"].items()])
    sentences = extract_sentences(report["report_text"])
    if not sentences:
        return {"error": "No checkable sentences found."}

    results = []
    hallucination_count = 0

    for sentence in sentences:
        nli_input = f"{source_text} [SEP] {sentence}"
        prediction = nli_pipeline(nli_input, truncation=True, max_length=512)
        label = prediction[0]["label"].lower()
        score = prediction[0]["score"]

        if "entail" in label:
            status = "supported"
        elif "contradict" in label:
            status = "hallucinated"
            hallucination_count += 1
        else:
            status = "unverified"
            if score > 0.80:
                hallucination_count += 0.5

        results.append({"sentence": sentence, "status": status, "confidence": round(score, 4)})

    total = len(sentences)
    hallucination_rate = round(hallucination_count / total, 4) if total > 0 else 0

    return {
        "report_type": report["report_type"],
        "total_claims": total,
        "hallucination_rate": hallucination_rate,
        "safety_score": round(1 - hallucination_rate, 4),
        "claim_results": results
    }


def evaluate_report(report):
    reference = " ".join([f"{k} is {v}." for k, v in report["patient_data"].items()])
    hypothesis = report["report_text"]

    P, R, F1 = bert_score([hypothesis], [reference], lang="en", verbose=False)
    scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
    rouge_scores = scorer.score(reference, hypothesis)

    return {
        "bertscore": {
            "precision": round(P[0].item(), 4),
            "recall":    round(R[0].item(), 4),
            "f1":        round(F1[0].item(), 4)
        },
        "rouge": {
            "rouge1": round(rouge_scores["rouge1"].fmeasure, 4),
            "rouge2": round(rouge_scores["rouge2"].fmeasure, 4),
            "rougeL": round(rouge_scores["rougeL"].fmeasure, 4)
        }
    }


def create_fhir_report(report, hallucination_result):
    report_type_codes = {
        "radiology": {"code": "18748-4", "display": "Diagnostic imaging study"},
        "discharge": {"code": "18842-5", "display": "Discharge summary"},
        "lab":       {"code": "11502-2", "display": "Laboratory report"}
    }
    code_info = report_type_codes.get(report["report_type"], {"code": "unknown", "display": "Clinical Report"})

    return {
        "resourceType": "DiagnosticReport",
        "id": str(uuid.uuid4()),
        "status": "final",
        "category": [{"coding": [{"system": "http://terminology.hl7.org/CodeSystem/v2-0074",
                                   "code": code_info["code"],
                                   "display": code_info["display"]}]}],
        "code": {"text": f"{report['report_type'].capitalize()} Report"},
        "subject": {"reference": f"Patient/{str(uuid.uuid4())}",
                    "display": report["patient_data"].get("name", "Unknown")},
        "effectiveDateTime": report["generated_at"],
        "issued": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ"),
        "conclusion": report["report_text"],
        "extension": [
            {"url": "https://medical-ai-portfolio.dev/fhir/hallucination-score",
             "valueDecimal": hallucination_result.get("hallucination_rate", 0)},
            {"url": "https://medical-ai-portfolio.dev/fhir/safety-score",
             "valueDecimal": hallucination_result.get("safety_score", 1)},
            {"url": "https://medical-ai-portfolio.dev/fhir/total-claims-checked",
             "valueInteger": hallucination_result.get("total_claims", 0)}
        ]
    }


def export_pdf(report, hallucination_result, eval_result):
    output_path = f"/tmp/{report['report_type']}_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
    doc = SimpleDocTemplate(output_path, pagesize=A4,
                            rightMargin=0.75*inch, leftMargin=0.75*inch,
                            topMargin=0.75*inch, bottomMargin=0.75*inch)
    styles = getSampleStyleSheet()

    title_style   = ParagraphStyle("title",   parent=styles["Title"],   fontSize=16, spaceAfter=6,  alignment=TA_CENTER)
    subtitle_style= ParagraphStyle("subtitle",parent=styles["Normal"],  fontSize=9,  spaceAfter=12, alignment=TA_CENTER, textColor=colors.grey)
    section_style = ParagraphStyle("section", parent=styles["Heading2"],fontSize=11, spaceBefore=12,spaceAfter=4, textColor=colors.HexColor("#1a1a2e"))
    body_style    = ParagraphStyle("body",    parent=styles["Normal"],  fontSize=9,  spaceAfter=6,  leading=14)
    label_style   = ParagraphStyle("label",   parent=styles["Normal"],  fontSize=9,  textColor=colors.grey)

    story = []
    story.append(Paragraph("Medical Report Generator", title_style))
    story.append(Paragraph(f"Medical AI Portfolio — Project 4 | Generated: {report['generated_at']}", subtitle_style))
    story.append(HRFlowable(width="100%", thickness=1, color=colors.HexColor("#1a1a2e")))
    story.append(Spacer(1, 12))
    story.append(Paragraph(f"{report['report_type'].upper()} REPORT", section_style))
    story.append(Spacer(1, 6))

    patient_table_data = [[Paragraph(f"<b>{k.replace('_',' ').title()}</b>", label_style),
                            Paragraph(str(v), body_style)]
                           for k, v in report["patient_data"].items()]
    patient_table = Table(patient_table_data, colWidths=[1.8*inch, 4.5*inch])
    patient_table.setStyle(TableStyle([
        ("BACKGROUND", (0,0),(0,-1), colors.HexColor("#f0f0f0")),
        ("GRID",       (0,0),(-1,-1),0.5, colors.lightgrey),
        ("VALIGN",     (0,0),(-1,-1),"TOP"),
        ("PADDING",    (0,0),(-1,-1),6),
    ]))
    story.append(Paragraph("Patient Information", section_style))
    story.append(patient_table)
    story.append(Spacer(1, 12))

    story.append(Paragraph("Generated Report", section_style))
    story.append(HRFlowable(width="100%", thickness=0.5, color=colors.lightgrey))
    story.append(Spacer(1, 6))
    for line in report["report_text"].split("\n"):
        if line.strip():
            if line.strip().isupper() or line.strip().endswith(":"):
                story.append(Paragraph(f"<b>{line.strip()}</b>", body_style))
            else:
                story.append(Paragraph(line.strip(), body_style))
    story.append(Spacer(1, 12))

    story.append(Paragraph("Quality Assessment", section_style))
    quality_data = [
        ["Metric", "Value"],
        ["Total Claims Checked",  str(hallucination_result.get("total_claims", 0))],
        ["Hallucination Rate",    str(hallucination_result.get("hallucination_rate", 0))],
        ["Safety Score",          str(hallucination_result.get("safety_score", 0))],
        ["BERTScore F1",          str(eval_result["bertscore"]["f1"])],
        ["ROUGE-1",               str(eval_result["rouge"]["rouge1"])],
        ["ROUGE-2",               str(eval_result["rouge"]["rouge2"])],
        ["ROUGE-L",               str(eval_result["rouge"]["rougeL"])],
    ]
    quality_table = Table(quality_data, colWidths=[2.5*inch, 2*inch])
    quality_table.setStyle(TableStyle([
        ("BACKGROUND", (0,0),(-1,0), colors.HexColor("#1a1a2e")),
        ("TEXTCOLOR",  (0,0),(-1,0), colors.white),
        ("FONTNAME",   (0,0),(-1,0), "Helvetica-Bold"),
        ("GRID",       (0,0),(-1,-1),0.5, colors.lightgrey),
        ("BACKGROUND", (0,1),(-1,-1),colors.HexColor("#f9f9f9")),
        ("PADDING",    (0,0),(-1,-1),6),
    ]))
    story.append(quality_table)
    story.append(Spacer(1, 12))
    story.append(HRFlowable(width="100%", thickness=0.5, color=colors.lightgrey))
    story.append(Paragraph("Generated by Medical Report Generator | Moumita Roy | Medical AI Portfolio Project 4", subtitle_style))
    doc.build(story)
    return output_path


# ── Pipeline ──────────────────────────────────────────────────
def run_pipeline(name, age, sex, chief_complaint, vitals, history, imaging, labs, report_type):
    patient_data = {
        "name": name, "age": age, "sex": sex,
        "chief_complaint": chief_complaint, "vitals": vitals,
        "history": history, "imaging": imaging, "labs": labs
    }
    patient_data = {k: v for k, v in patient_data.items() if v and str(v).strip()}

    report       = generate_report(patient_data, report_type)
    h_result     = check_hallucination(report)
    e_result     = evaluate_report(report)
    fhir         = create_fhir_report(report, h_result)
    pdf_path     = export_pdf(report, h_result, e_result)

    claim_breakdown = "\n".join([
        f"[{c['status'].upper():12}] ({c['confidence']}) {c['sentence'][:90]}..."
        for c in h_result.get("claim_results", [])
    ])

    hallucination_summary = f"""Total Claims Checked : {h_result['total_claims']}
Hallucination Rate   : {h_result['hallucination_rate']}
Safety Score         : {h_result['safety_score']}

Claim-level Breakdown:
{claim_breakdown}"""

    metrics_summary = f"""BERTScore  —  P: {e_result['bertscore']['precision']}  R: {e_result['bertscore']['recall']}  F1: {e_result['bertscore']['f1']}
ROUGE-1    :  {e_result['rouge']['rouge1']}
ROUGE-2    :  {e_result['rouge']['rouge2']}
ROUGE-L    :  {e_result['rouge']['rougeL']}"""

    return (
        report["report_text"],
        hallucination_summary,
        metrics_summary,
        json.dumps(fhir, indent=2),
        pdf_path
    )


# ── Gradio UI ─────────────────────────────────────────────────
with gr.Blocks(title="Medical Report Generator") as demo:
    gr.Markdown("# Medical Report Generator")
    gr.Markdown("Medical AI Portfolio — Project 4 | Moumita Roy | [GitHub](https://github.com/moumitaroy19/medical-report-generator)")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Patient Information")
            name            = gr.Textbox(label="Full Name",        value="John Doe")
            age             = gr.Textbox(label="Age",              value="58")
            sex             = gr.Textbox(label="Sex",              value="Male")
            chief_complaint = gr.Textbox(label="Chief Complaint",  value="Chest pain and shortness of breath for 2 days")
            vitals          = gr.Textbox(label="Vitals",           value="BP 145/90, HR 88, RR 18, Temp 37.1C, SpO2 96%")
            history         = gr.Textbox(label="Medical History",  value="Hypertension, Type 2 Diabetes, smoker for 20 years")
            imaging         = gr.Textbox(label="Imaging",          value="Chest X-ray ordered, mild cardiomegaly noted")
            labs            = gr.Textbox(label="Lab Results",      value="WBC 11.2, HGB 13.4, Troponin 0.02, BNP 210")
            report_type     = gr.Dropdown(choices=["radiology", "discharge", "lab"],
                                          value="radiology", label="Report Type")
            submit_btn      = gr.Button("Generate Report", variant="primary")

        with gr.Column():
            gr.Markdown("### Output")
            report_output        = gr.Textbox(label="Generated Report",       lines=12)
            hallucination_output = gr.Textbox(label="Hallucination Analysis", lines=10)
            metrics_output       = gr.Textbox(label="Evaluation Metrics",     lines=6)
            fhir_output          = gr.Textbox(label="FHIR R4 JSON",           lines=10)
            pdf_output           = gr.File(label="Download PDF Report")

    submit_btn.click(
        fn=run_pipeline,
        inputs=[name, age, sex, chief_complaint, vitals, history, imaging, labs, report_type],
        outputs=[report_output, hallucination_output, metrics_output, fhir_output, pdf_output]
    )

demo.launch()