import fitz  # PyMuPDF
import json
import os
import io
import re
import tempfile
from docx import Document
from docx.shared import Inches, Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.enum.table import WD_TABLE_ALIGNMENT
from docx.oxml.ns import qn
from docx.oxml import OxmlElement
import gradio as gr
from datetime import datetime


# ─────────────────────────────────────────────
# 1. PDF EXTRACTION
# ─────────────────────────────────────────────

def extract_pdf_content(pdf_path):
    """Extract text and images from a PDF, grouped by page."""
    doc = fitz.open(pdf_path)
    pages_text = []
    images = []  # list of (page_num, image_bytes, img_index)

    for page_num, page in enumerate(doc):
        text = page.get_text("text")
        pages_text.append(text)

        # Extract embedded images
        img_list = page.get_images(full=True)
        for img_idx, img_info in enumerate(img_list):
            xref = img_info[0]
            try:
                base_image = doc.extract_image(xref)
                img_bytes = base_image["image"]
                img_ext = base_image["ext"]
                if len(img_bytes) > 5000:  # Skip tiny icons
                    images.append({
                        "page": page_num,
                        "bytes": img_bytes,
                        "ext": img_ext,
                        "xref": xref
                    })
            except Exception:
                continue

    full_text = "\n".join(pages_text)
    doc.close()
    return full_text, images


def parse_thermal_data(thermal_text):
    """
    Parse Bosch GTC 400C thermal report into structured readings.
    Returns a formatted summary string for groq.
    """
    readings = []
    pattern = re.compile(
        r'Thermal image\s*:\s*(\S+\.JPG).*?'
        r'Hotspot\s*:\s*([\d.]+)\s*°C.*?'
        r'Coldspot\s*:\s*([\d.]+)\s*°C',
        re.DOTALL | re.IGNORECASE
    )
    for i, match in enumerate(pattern.finditer(thermal_text), 1):
        filename, hotspot, coldspot = match.groups()
        delta = round(float(hotspot) - float(coldspot), 1)
        if delta >= 5:
            level = "HIGH moisture"
        elif delta >= 3:
            level = "MODERATE moisture"
        else:
            level = "LOW/trace moisture"
        readings.append(
            f"Reading {i} ({filename}): Hotspot={hotspot}°C, Coldspot={coldspot}°C, Delta={delta}°C → {level}"
        )
    if not readings:
        return thermal_text[:8000]
    summary = f"THERMAL REPORT — {len(readings)} readings, Bosch GTC 400C, 27/09/2022\n"
    summary += "Ambient ~23°C. Coldspots below ambient = moisture confirmed.\n\n"
    summary += "\n".join(readings)
    return summary


def parse_impacted_areas_from_text(text):
    """
    Parse the structured impacted areas from the UrbanRoof report.
    Returns a list of dicts with area name, negative description, positive description.
    """
    areas = []

    # Extract Summary Table section
    summary_match = re.search(
        r'SUMMARY TABLE(.*?)(?:Appendix|$)', text, re.DOTALL | re.IGNORECASE
    )
    summary_text = summary_match.group(1) if summary_match else ""

    # Parse numbered rows from summary table
    # Pattern: number followed by description text
    rows = re.findall(
        r'(\d+)\s+(Observed[^\n]+(?:\n(?!\d).*)*)',
        summary_text
    )

    for row_num, row_text in rows:
        row_text = row_text.strip().replace('\n', ' ')
        areas.append({
            "point": row_num,
            "description": row_text
        })

    # Also extract impacted area sections for richer data
    area_sections = re.findall(
        r'Impacted Area (\d+)\s*Negative side Description\s*([^\n]+(?:\n(?!Positive|Impacted).*)*)'
        r'.*?Positive side Description\s*([^\n]+(?:\n(?!Negative|Impacted).*)*)',
        text, re.DOTALL
    )

    area_details = []
    for match in area_sections:
        area_num, neg_desc, pos_desc = match
        area_details.append({
            "area_num": area_num.strip(),
            "negative": neg_desc.strip().replace('\n', ' '),
            "positive": pos_desc.strip().replace('\n', ' ')
        })

    return areas, area_details


# ─────────────────────────────────────────────
# 2. groq AI ANALYSIS
# ─────────────────────────────────────────────

DDR_PROMPT = """You are an expert civil engineer analyzing a property inspection report and thermal imaging report from UrbanRoof.

INSPECTION REPORT TEXT:
{inspection_text}

{thermal_section}

THERMAL DATA INTERPRETATION GUIDE:
- The thermal report uses a Bosch GTC 400C infrared camera (emissivity 0.94)
- Cold spots (blue/cyan in thermal images) = moisture/dampness presence
- A coldspot significantly below ambient (23°C) confirms active moisture
- Temperature delta (hotspot - coldspot) > 4°C indicates significant moisture
- Readings taken same day as inspection: 27/09/2022

Based ONLY on the information in the reports above, generate a structured Detailed Diagnostic Report (DDR).

STRICT RULES:
- Do NOT invent any facts not present in the documents
- If information is missing → write "Not Available"  
- If thermal data conflicts with visual inspection → note the conflict explicitly
- Use thermal readings to CONFIRM or QUALIFY the severity of each impacted area
- Return ONLY valid JSON with no markdown fences, no preamble, no explanation

Return this exact JSON structure:
{{
  "property_summary": {{
    "property_type": "...",
    "floors": "...",
    "inspection_date": "...",
    "inspected_by": "...",
    "overall_score": "...",
    "flagged_items": "...",
    "previous_structural_audit": "...",
    "previous_repair_work": "..."
  }},
  "executive_summary": "2-3 sentence summary combining inspection findings with thermal confirmation",
  "area_wise_observations": [
    {{
      "point_no": "1",
      "area": "area name (e.g. Hall, Bedroom)",
      "impacted_side_observation": "what was observed on the negative/impacted side",
      "source_side_observation": "what was found on the positive/exposed side (root cause location)",
      "thermal_finding": "thermal camera reading for this area: hotspot temp, coldspot temp, delta, and what it indicates. Write 'Not Available' if no thermal reading correlates to this area.",
      "severity": "High/Medium/Low — use thermal delta to inform severity: >5°C delta = High, 3-5°C = Medium, <3°C = Low"
    }}
  ],
  "probable_root_cause": "Detailed root cause analysis combining visual inspection and thermal evidence",
  "severity_assessment": {{
    "overall": "High/Medium/Low",
    "reasoning": "Why this overall severity — cite both visual and thermal evidence",
    "breakdown": [
      {{"issue": "issue name", "severity": "High/Medium/Low", "basis": "combined visual + thermal basis"}}
    ]
  }},
  "recommended_actions": [
    {{
      "priority": "Immediate/Short-term/Long-term",
      "action": "specific action to take",
      "area": "which area this applies to"
    }}
  ],
  "checklist_findings": {{
    "wc_checklist": "summary of WC checklist findings",
    "external_wall_checklist": "summary of external wall findings",
    "structural_condition": "summary of structural condition"
  }},
  "thermal_summary": "Overall interpretation of the 30 thermal readings: moisture extent, worst affected areas by thermal delta, whether thermal confirms or contradicts visual findings",
  "additional_notes": "Any other important observations",
  "missing_or_unclear": ["list any information that was missing or unclear in the reports"]
}}"""


def analyze_with_groq(api_key, inspection_text, thermal_text=""):
    from groq import Groq

    client = Groq(api_key=api_key)

    thermal_section = ""
    if thermal_text and thermal_text.strip():
        parsed_thermal = parse_thermal_data(thermal_text)
        thermal_section = f"THERMAL REPORT TEXT:\n{parsed_thermal[:10000]}"

    prompt = DDR_PROMPT.format(
        inspection_text=inspection_text[:15000],
        thermal_section=thermal_section
    )

    response = client.chat.completions.create(
        model="llama-3.3-70b-versatile",
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"}
    )

    raw = response.choices[0].message.content.strip()

    try:
        return json.loads(raw)
    except json.JSONDecodeError as e:
        json_match = re.search(r'\{.*\}', raw, re.DOTALL)
        if json_match:
            return json.loads(json_match.group())
        raise ValueError(f"Could not parse response as JSON: {e}\n\nRaw: {raw[:500]}")
# ─────────────────────────────────────────────
# 3. DOCX REPORT BUILDER
# ─────────────────────────────────────────────

def set_cell_bg(cell, hex_color):
    """Set table cell background color."""
    tc = cell._tc
    tcPr = tc.get_or_add_tcPr()
    shd = OxmlElement('w:shd')
    shd.set(qn('w:val'), 'clear')
    shd.set(qn('w:color'), 'auto')
    shd.set(qn('w:fill'), hex_color)
    tcPr.append(shd)


def add_heading(doc, text, level=1, color=None):
    """Add a styled heading."""
    heading = doc.add_heading(text, level=level)
    if color:
        for run in heading.runs:
            run.font.color.rgb = RGBColor(*color)
    return heading


def add_colored_paragraph(doc, text, bold=False, color=None, size=None):
    """Add a paragraph with optional styling."""
    para = doc.add_paragraph()
    run = para.add_run(text)
    run.bold = bold
    if color:
        run.font.color.rgb = RGBColor(*color)
    if size:
        run.font.size = Pt(size)
    return para


def build_ddr_document(ddr_data, images, output_path):
    """Build the full DDR Word document from parsed data and images."""
    doc = Document()

    # Page margins
    for section in doc.sections:
        section.top_margin = Inches(1)
        section.bottom_margin = Inches(1)
        section.left_margin = Inches(1.2)
        section.right_margin = Inches(1.2)

    # ── COVER / HEADER ──
    title = doc.add_heading('DETAILED DIAGNOSTIC REPORT', 0)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title.runs[0].font.color.rgb = RGBColor(0x8B, 0x1A, 0x1A)

    subtitle = doc.add_paragraph('Property Inspection Analysis — UrbanRoof')
    subtitle.alignment = WD_ALIGN_PARAGRAPH.CENTER
    subtitle.runs[0].bold = True
    subtitle.runs[0].font.size = Pt(13)

    doc.add_paragraph(f'Report Generated: {datetime.now().strftime("%d %B %Y")}').alignment = WD_ALIGN_PARAGRAPH.CENTER
    doc.add_paragraph('')

    # ── SECTION 1: PROPERTY SUMMARY ──
    prop = ddr_data.get('property_summary', {})
    add_heading(doc, '1. Property Summary', level=1, color=(0x8B, 0x1A, 0x1A))

    summary_table = doc.add_table(rows=0, cols=2)
    summary_table.style = 'Table Grid'

    fields = [
        ('Property Type', prop.get('property_type', 'Not Available')),
        ('Number of Floors', prop.get('floors', 'Not Available')),
        ('Inspection Date', prop.get('inspection_date', 'Not Available')),
        ('Inspected By', prop.get('inspected_by', 'Not Available')),
        ('Overall Score', prop.get('overall_score', 'Not Available')),
        ('Flagged Items', prop.get('flagged_items', 'Not Available')),
        ('Previous Structural Audit', prop.get('previous_structural_audit', 'Not Available')),
        ('Previous Repair Work', prop.get('previous_repair_work', 'Not Available')),
    ]

    for label, value in fields:
        row = summary_table.add_row()
        row.cells[0].text = label
        row.cells[0].paragraphs[0].runs[0].bold = True
        set_cell_bg(row.cells[0], 'F2F2F2')
        row.cells[1].text = str(value)

    doc.add_paragraph('')

    # ── SECTION 2: EXECUTIVE SUMMARY ──
    add_heading(doc, '2. Executive Summary', level=1, color=(0x8B, 0x1A, 0x1A))
    doc.add_paragraph(ddr_data.get('executive_summary', 'Not Available'))
    doc.add_paragraph('')

    # ── SECTION 3: AREA-WISE OBSERVATIONS ──
    add_heading(doc, '3. Area-Wise Observations', level=1, color=(0x8B, 0x1A, 0x1A))

    observations = ddr_data.get('area_wise_observations', [])
    if observations:
        obs_table = doc.add_table(rows=1, cols=6)
        obs_table.style = 'Table Grid'

        # Header row
        headers = ['Point', 'Area', 'Observed Damage (Impacted Side)', 'Root Source (Exposed Side)', 'Thermal Finding', 'Severity']
        hdr_row = obs_table.rows[0]
        for i, hdr in enumerate(headers):
            cell = hdr_row.cells[i]
            cell.text = hdr
            cell.paragraphs[0].runs[0].bold = True
            set_cell_bg(cell, '8B1A1A')
            cell.paragraphs[0].runs[0].font.color.rgb = RGBColor(0xFF, 0xFF, 0xFF)

        # Data rows
        severity_colors = {'High': 'FFE0E0', 'Medium': 'FFF3CD', 'Low': 'E8F5E9'}
        for obs in observations:
            row = obs_table.add_row()
            row.cells[0].text = str(obs.get('point_no', ''))
            row.cells[1].text = obs.get('area', '')
            row.cells[2].text = obs.get('impacted_side_observation', '')
            row.cells[3].text = obs.get('source_side_observation', '')
            row.cells[4].text = obs.get('thermal_finding', 'Not Available')
            severity = obs.get('severity', 'Medium')
            row.cells[5].text = severity
            bg = severity_colors.get(severity, 'FFFFFF')
            set_cell_bg(row.cells[5], bg)

    # Add first few inspection photos after observations
    doc.add_paragraph('')
    if images:
        add_heading(doc, 'Inspection Photographs — Impacted Areas', level=2)
        # Add photos in groups of 2 per row
        photo_count = 0
        for img_data in images[:12]:  # First 12 photos (impacted area shots)
            try:
                img_stream = io.BytesIO(img_data['bytes'])
                doc.add_picture(img_stream, width=Inches(2.8))
                photo_count += 1
            except Exception:
                continue
        doc.add_paragraph('')

    # ── SECTION 4: ROOT CAUSE ANALYSIS ──
    add_heading(doc, '4. Probable Root Cause', level=1, color=(0x8B, 0x1A, 0x1A))
    doc.add_paragraph(ddr_data.get('probable_root_cause', 'Not Available'))
    doc.add_paragraph('')

    # ── SECTION 5: SEVERITY ASSESSMENT ──
    add_heading(doc, '5. Severity Assessment', level=1, color=(0x8B, 0x1A, 0x1A))

    sev = ddr_data.get('severity_assessment', {})
    overall_sev = sev.get('overall', 'Not Available')

    # Overall severity box
    sev_para = doc.add_paragraph()
    sev_run = sev_para.add_run(f'Overall Severity: {overall_sev}')
    sev_run.bold = True
    sev_run.font.size = Pt(13)
    color_map = {'High': RGBColor(0xCC, 0x00, 0x00), 'Medium': RGBColor(0xFF, 0x8C, 0x00), 'Low': RGBColor(0x00, 0x80, 0x00)}
    sev_run.font.color.rgb = color_map.get(overall_sev, RGBColor(0, 0, 0))

    doc.add_paragraph(sev.get('reasoning', ''))

    breakdown = sev.get('breakdown', [])
    if breakdown:
        sev_table = doc.add_table(rows=1, cols=3)
        sev_table.style = 'Table Grid'
        hdr_cells = sev_table.rows[0].cells
        for i, h in enumerate(['Issue', 'Severity', 'Basis']):
            hdr_cells[i].text = h
            hdr_cells[i].paragraphs[0].runs[0].bold = True
            set_cell_bg(hdr_cells[i], 'F2F2F2')

        for item in breakdown:
            row = sev_table.add_row()
            row.cells[0].text = item.get('issue', '')
            row.cells[1].text = item.get('severity', '')
            row.cells[2].text = item.get('basis', '')

    doc.add_paragraph('')

    # ── SECTION 6: RECOMMENDED ACTIONS ──
    add_heading(doc, '6. Recommended Actions', level=1, color=(0x8B, 0x1A, 0x1A))

    actions = ddr_data.get('recommended_actions', [])
    if actions:
        act_table = doc.add_table(rows=1, cols=3)
        act_table.style = 'Table Grid'
        priority_cells = act_table.rows[0].cells
        for i, h in enumerate(['Priority', 'Recommended Action', 'Area']):
            priority_cells[i].text = h
            priority_cells[i].paragraphs[0].runs[0].bold = True
            set_cell_bg(priority_cells[i], '8B1A1A')
            priority_cells[i].paragraphs[0].runs[0].font.color.rgb = RGBColor(0xFF, 0xFF, 0xFF)

        priority_colors = {'Immediate': 'FFE0E0', 'Short-term': 'FFF3CD', 'Long-term': 'E8F5E9'}
        for action in actions:
            row = act_table.add_row()
            priority = action.get('priority', '')
            row.cells[0].text = priority
            row.cells[1].text = action.get('action', '')
            row.cells[2].text = action.get('area', '')
            set_cell_bg(row.cells[0], priority_colors.get(priority, 'FFFFFF'))

    doc.add_paragraph('')

    # ── SECTION 7: CHECKLIST FINDINGS ──
    add_heading(doc, '7. Checklist Findings', level=1, color=(0x8B, 0x1A, 0x1A))
    checklist = ddr_data.get('checklist_findings', {})

    for key, label in [('wc_checklist', 'WC Checklist'), ('external_wall_checklist', 'External Wall'), ('structural_condition', 'Structural Condition')]:
        if checklist.get(key):
            add_heading(doc, label, level=2)
            doc.add_paragraph(checklist[key])

    # Additional source-side photos
    if len(images) > 12:
        doc.add_paragraph('')
        add_heading(doc, 'Inspection Photographs — Source Areas', level=2)
        for img_data in images[12:24]:
            try:
                img_stream = io.BytesIO(img_data['bytes'])
                doc.add_picture(img_stream, width=Inches(2.8))
            except Exception:
                continue

    # ── THERMAL SUMMARY ──
    thermal_summary = ddr_data.get('thermal_summary', '')
    if thermal_summary and thermal_summary != 'Not Available':
        doc.add_paragraph('')
        add_heading(doc, '8. Thermal Imaging Summary', level=1, color=(0x8B, 0x1A, 0x1A))
        doc.add_paragraph(thermal_summary)

    # ── SECTION 8/9: ADDITIONAL NOTES & MISSING INFO ──
    doc.add_paragraph('')
    add_heading(doc, '9. Additional Notes', level=1, color=(0x8B, 0x1A, 0x1A))
    doc.add_paragraph(ddr_data.get('additional_notes', 'None'))

    missing = ddr_data.get('missing_or_unclear', [])
    if missing:
        add_heading(doc, 'Missing or Unclear Information', level=2)
        for item in missing:
            para = doc.add_paragraph(style='List Bullet')
            para.add_run(item)

    doc.save(output_path)
    return output_path


# ─────────────────────────────────────────────
# 4. GRADIO UI
# ─────────────────────────────────────────────

def generate_ddr(api_key, inspection_pdf, thermal_pdf):
    """Main pipeline: PDF → AI analysis → DDR Word document."""

    if not api_key or not api_key.strip():
        return None, "❌ Please enter your groq API key. Get one free at: https://aistudio.google.com/app/apikey"

    if inspection_pdf is None:
        return None, "❌ Please upload an Inspection Report PDF."

    try:
        # Step 1: Extract inspection PDF
        yield None, "📄 Step 1/4: Extracting inspection report..."
        inspection_text, images = extract_pdf_content(inspection_pdf)

        if len(inspection_text.strip()) < 100:
            return None, "❌ Could not extract text from the inspection PDF. It may be a scanned image without OCR."

        # Step 2: Extract thermal PDF (optional)
        thermal_text = ""
        if thermal_pdf is not None:
            yield None, "📄 Step 2/4: Extracting thermal report..."
            thermal_text, thermal_images = extract_pdf_content(thermal_pdf)
            images = images + thermal_images
        else:
            yield None, "⚠️  No thermal report provided — proceeding with inspection report only."

        # Step 3: AI analysis
        yield None, f"🤖 Step 3/4: Analyzing with groq AI... ({len(images)} photos found)"
        ddr_data = analyze_with_groq(api_key.strip(), inspection_text, thermal_text)

        # Step 4: Build document
        yield None, "📝 Step 4/4: Building DDR Word document..."
        output_path = tempfile.mktemp(suffix=".docx")
        build_ddr_document(ddr_data, images, output_path)

        areas_count = len(ddr_data.get('area_wise_observations', []))
        actions_count = len(ddr_data.get('recommended_actions', []))
        overall_sev = ddr_data.get('severity_assessment', {}).get('overall', 'N/A')

        summary = f"""✅ DDR Generated Successfully!

📊 Summary:
• {areas_count} impacted areas documented
• {actions_count} recommended actions
• Overall severity: {overall_sev}
• {len(images)} photos embedded

Download your report below ↓"""

        yield output_path, summary

    except Exception as e:
        import traceback
        yield None, f"❌ Error: {str(e)}\n\nDetails:\n{traceback.format_exc()}"


with gr.Blocks(title="DDR Generator — UrbanRoof", theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # 🏠 Detailed Diagnostic Report Generator
    **UrbanRoof Property Inspection → Professional DDR Word Document**
    
    Upload your inspection report (and optionally a thermal report), and the AI will generate a structured DDR with observations, root cause analysis, severity assessment, and recommended actions.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            api_key_input = gr.Textbox(
                label="🔑 groq API Key",
                placeholder="Paste your groq API key here (free at aistudio.google.com)",
                type="password"
            )
            inspection_input = gr.File(
                label="📋 Inspection Report PDF (Required)",
                file_types=[".pdf"]
            )
            thermal_input = gr.File(
                label="🌡️ Thermal Report PDF (Optional)",
                file_types=[".pdf"]
            )
            generate_btn = gr.Button("🚀 Generate DDR Report", variant="primary", size="lg")

        with gr.Column(scale=1):
            status_output = gr.Textbox(
                label="Status",
                lines=10,
                interactive=False,
                value="Upload your PDFs and click Generate to begin."
            )
            file_output = gr.File(label="📥 Download DDR Report (.docx)")

    generate_btn.click(
        fn=generate_ddr,
        inputs=[api_key_input, inspection_input, thermal_input],
        outputs=[file_output, status_output]
    )

    gr.Markdown("""
    ---
    **How it works:** PyMuPDF extracts text + photos → groq analyzes and structures the data → python-docx builds the Word report
    
    **Note:** The thermal report is optional. If not provided, the DDR will be generated from the inspection report alone.
    """)


if __name__ == "__main__":
    app.launch(share=True)