import os
import base64
import math
from datetime import datetime, timezone
from zoneinfo import ZoneInfo
from io import BytesIO

import requests
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import cm
from reportlab.lib.utils import ImageReader
from reportlab.lib import colors
from PIL import Image, ImageOps
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from pydantic import BaseModel

# Configuration from environment variables
BASE_URL = os.getenv("BASE_URL", "https://rakshitjan-cps-b2c.hf.space")
ADMIN_EMAIL = os.getenv("ADMIN_EMAIL", "admin123@example.com")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "SecurePass123")
AES_KEY_STR = os.getenv("AES_KEY_STR", "7E892875A52C59A3B588306B13C31FBD")
AES_IV_STR = os.getenv("AES_IV_STR", "XYE45DKJ0967GFAZ")
GEOAPIFY_API_KEY = os.getenv("GEOAPIFY_API_KEY", "ceaea8a8a8f14f46b73221384eb4b5f0")
LOGO_PATH = os.getenv("LOGO_PATH", "logo.png")

app = FastAPI(title="Case Report Generator", version="1.0.0")

# AES Encryption/Decryption functions
def aes_encrypt(plaintext: str, key_str: str, iv_str: str) -> str:
    key = key_str.encode("utf-8")
    iv  = iv_str.encode("utf-8")
    cipher = AES.new(key, AES.MODE_CBC, iv)
    ct = cipher.encrypt(pad(plaintext.encode("utf-8"), AES.block_size))
    return base64.b64encode(ct).decode("utf-8")

def aes_decrypt(enc_b64: str, key_str: str, iv_str: str) -> str:
    try:
        key = key_str.encode("utf-8")
        iv  = iv_str.encode("utf-8")
        cipher = AES.new(key, AES.MODE_CBC, iv)
        pt = unpad(cipher.decrypt(base64.b64decode(enc_b64)), AES.block_size)
        return pt.decode("utf-8")
    except Exception:
        return enc_b64

# API Client functions
session = requests.Session()
session.headers.update({"Content-Type": "application/json"})

def admin_login():
    payload = {"email": aes_encrypt(ADMIN_EMAIL, AES_KEY_STR, AES_IV_STR),
               "password": ADMIN_PASSWORD}
    r = session.post(f"{BASE_URL}/api/admin/login", json=payload, timeout=30)
    data = r.json()
    if r.status_code != 200 or "token" not in data:
        raise RuntimeError(f"Login failed: {data}")
    session.headers.update({"Authorization": f"Bearer {data['token']}"})
    print("✅ Admin Logged in")
    return data["token"]

def get_case(case_id: str):
    r = session.get(f"{BASE_URL}/api/cases/{case_id}", timeout=30)
    r.raise_for_status()
    return r.json()

def get_files(case_id: str):
    r = session.get(f"{BASE_URL}/api/cases/{case_id}/files", timeout=60)
    r.raise_for_status()
    return r.json().get("files", [])

def download_file(file_id: str):
    r = session.get(f"{BASE_URL}/api/files/{file_id}", timeout=60)
    if r.status_code == 200:
        return r.content
    raise RuntimeError(f"File download failed for {file_id}: {r.status_code}")

# Geoapify functions
def forward_geocode(address: str):
    if not address: return (None, None)
    url = f"https://api.geoapify.com/v1/geocode/search?text={requests.utils.quote(address)}&apiKey={GEOAPIFY_API_KEY}"
    res = requests.get(url, timeout=20).json()
    if res.get("features"):
        p = res["features"][0]["properties"]
        return p.get("lat"), p.get("lon")
    return (None, None)

def reverse_geocode(lat: float, lng: float):
    if lat is None or lng is None: return "Unknown Address"
    url = f"https://api.geoapify.com/v1/geocode/reverse?lat={lat}&lon={lng}&apiKey={GEOAPIFY_API_KEY}"
    res = requests.get(url, timeout=20).json()
    return res.get("features", [{}])[0].get("properties", {}).get("formatted", "Unknown Address")

def static_map_with_path(capture_lat, capture_lng, case_lat=None, case_lng=None, zoom: int = 17, size="1800x600"):
    """Map shows capture pin (red), case pin (blue), path between them."""
    if capture_lat is None or capture_lng is None:
        return None
    parts = [
        f"center=lonlat:{capture_lng},{capture_lat}",
        f"zoom={zoom}",
        f"size={size}",
        f"scaleFactor=2",
        f"marker=lonlat:{capture_lng},{capture_lat};color:%23ff0000;size:large"
    ]
    if case_lat is not None and case_lng is not None:
        parts.append(f"marker=lonlat:{case_lng},{case_lat};color:%23007bff;size:large")
        parts.append(f"path=color:%23ff0000;weight:4|lonlat:{case_lng},{case_lat}|lonlat:{capture_lng},{capture_lat}")

    url = "https://maps.geoapify.com/v1/staticmap?" + "&".join(parts) + f"&apiKey={GEOAPIFY_API_KEY}"
    r = requests.get(url, timeout=30)
    if r.status_code == 200:
        return Image.open(BytesIO(r.content)).convert("RGB")
    return None

def haversine_km(lat1, lon1, lat2, lon2):
    if None in (lat1, lon1, lat2, lon2): return None
    R = 6371.0  # km
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlambda = math.radians(lon2 - lon1)
    a = math.sin(dphi/2)**2 + math.cos(phi1)*math.cos(phi2)*math.sin(dlambda/2)**2
    return 2 * R * math.atan2(math.sqrt(a), math.sqrt(1-a))

# Utility functions
IST = ZoneInfo("Asia/Kolkata")

def fmt_ist(dt_str: str) -> str:
    """ISO string -> 'DD/MM/YYYY hh:mm:ss AM/PM IST'"""
    if not dt_str:
        return "N/A"
    try:
        if dt_str.endswith("Z"):
            dt = datetime.fromisoformat(dt_str.replace("Z", "+00:00"))
        else:
            dt = datetime.fromisoformat(dt_str)
        dt_ist = dt.astimezone(IST)
        return dt_ist.strftime("%d/%m/%Y %I:%M:%S %p IST")
    except Exception:
        return dt_str

def exported_on_str() -> str:
    now_ist = datetime.now(IST)
    return now_ist.strftime("%B %d, %Y")

def draw_wrapped_label_value(c, x, y, label, text, max_width_cm=16.0):
    """Draw 'Label: value...' with word wrapping; following lines indented."""
    max_w = max_width_cm * cm
    c.setFont("Helvetica", 10)
    label_w = c.stringWidth(f"{label}: ", "Helvetica", 10)
    words = (text or "N/A")
    words = str(words).split()
    first = True
    line = ""
    width = 0
    lines = []

    for w in words:
        w_w = c.stringWidth(w + " ", "Helvetica", 10)
        if first:
            if label_w + width + w_w < max_w:
                line += w + " "; width += w_w
            else:
                lines.append(f"{label}: {line.strip()}")
                first = False
                line, width = w + " ", w_w
        else:
            if width + w_w < max_w:
                line += w + " "; width += w_w
            else:
                lines.append(" " * int(label_w/6) + line.strip())
                line, width = w + " ", w_w

    if first:
        lines.append(f"{label}: {line.strip()}")
    else:
        if line:
            lines.append(" " * int(label_w/6) + line.strip())

    for ln in lines:
        c.drawString(x, y, ln); y -= 0.45*cm
    return y

DISPLAY_NAME = {
    "aadhaar_photo": "Photo of Aadhar",
    "customer_photo": "Customer's Photo",
    "residence_photo": "Residence Photo",
    "business_photo": "Business Address Photo"
}

def display_type(file_type: str) -> str:
    return DISPLAY_NAME.get(file_type, file_type or "Upload")

def make_thumb(img_bytes: bytes, max_w=860, max_h=640):
    img = Image.open(BytesIO(img_bytes)).convert("RGB")
    img.thumbnail((max_w, max_h))
    return ImageOps.expand(img, border=1, fill=(170,170,170))

# PDF Generation functions
W, H = A4
MARGIN = 1.2*cm

def draw_header(c, title):
    c.setFillColor(colors.HexColor("#0F61A8"))
    c.rect(0, H-2.4*cm, W, 2.4*cm, fill=1, stroke=0)

    title_x = MARGIN
    if os.path.exists(LOGO_PATH):
        try:
            c.drawImage(LOGO_PATH, MARGIN, H-2.2*cm,
                        width=2.4*cm, height=2.0*cm, preserveAspectRatio=True, mask='auto')
            title_x = MARGIN + 2.6*cm
        except:
            pass

    c.setFillColor(colors.white)
    c.setFont("Helvetica-Bold", 15)
    c.drawString(title_x, H-1.2*cm, title)

    c.setFont("Helvetica", 10)
    exp = f"Exported on {exported_on_str()}"
    c.drawRightString(W - MARGIN, H-1.0*cm, exp)
    c.setFillColor(colors.black)

def draw_footer(c, page):
    c.setFont("Helvetica", 9)
    c.drawRightString(W - MARGIN, 1.0*cm, f"Page {page}")

def section_title(c, text, x, y):
    c.setFont("Helvetica-Bold", 14)
    c.drawString(x, y, text)
    return y - 0.9*cm

def label_value(c, x, y, label, value):
    c.setFont("Helvetica-Bold", 10)
    c.drawString(x, y, f"{label}:")
    c.setFont("Helvetica", 10)
    c.drawString(x + 4.8*cm, y, str("N/A" if value in (None, "", []) else value))

def build_pdf(case: dict, files: list, out_path: str):
    c = canvas.Canvas(out_path, pagesize=A4)
    page = 1

    ad = case.get("demographic_details", {}).get("address_details", {})
    case_addr = f"{ad.get('residential_address','')}, {ad.get('city','')}, {ad.get('state','')} {ad.get('pincode','')}"
    case_lat, case_lng = forward_geocode(case_addr)

    title = f"Name : {case.get('case_applicant_name','')}"
    draw_header(c, title)
    y = H - 3.4*cm

    y = section_title(c, "Customer Details", MARGIN, y)
    label_value(c, MARGIN, y, "Customer Name", case.get("case_applicant_name")); y -= 0.6*cm
    label_value(c, MARGIN, y, "Customer ID", case.get("case_id")); y -= 0.6*cm
    label_value(c, MARGIN, y, "Phone Number", case.get("case_applicant_contact")); y -= 0.6*cm
    email_id = case.get("demographic_details",{}).get("contact_information",{}).get("email_id")
    label_value(c, MARGIN, y, "Email", email_id); y -= 0.6*cm
    y = draw_wrapped_label_value(c, MARGIN, y, "Customer Address", case_addr, max_width_cm=18.0); y -= 0.4*cm

    y = section_title(c, "Inspection Summary", MARGIN, y)
    c.setFont("Helvetica", 10)
    c.drawString(MARGIN, y, f"Request Started: {fmt_ist(case.get('created_at'))}"); y -= 0.5*cm
    c.drawString(MARGIN, y, f"Total Uploads: {len(files)} uploads"); y -= 0.5*cm
    c.drawString(MARGIN, y, f"Status: {case.get('status')}"); y -= 0.5*cm
    c.drawString(MARGIN, y, f"Priority: {case.get('priority')}"); y -= 0.7*cm

    if case_lat:
        smap = static_map_with_path(case_lat, case_lng, case_lat, case_lng, zoom=12)
        if smap:
            map_h = H/2 + 2*cm
            map_w = W - 2*MARGIN
            c.drawImage(
                ImageReader(smap),
                MARGIN,
                y - map_h,
                width=map_w,
                height=map_h,
                preserveAspectRatio=True,
                mask='auto'
            )
            y -= map_h + 0.8*cm

    draw_footer(c, page)
    c.showPage(); page += 1

    draw_header(c, title)
    y = H - 3.4*cm
    c.setFont("Helvetica-Bold", 14)
    c.drawString(MARGIN, y, "Image Thumbnails")
    y -= 0.8*cm

    x0, y0 = MARGIN, y
    col_w = (W - 2*MARGIN) / 3
    row_h = 6.1*cm

    for i, f in enumerate(files[:6]):
        img = make_thumb(f["_bin"])
        col = i % 3
        row = i // 3
        xi = x0 + col * col_w + 0.2*cm
        yi = y0 - row * row_h - 4.6*cm
        c.drawImage(ImageReader(img), xi, yi,
                    width=col_w-0.6*cm, height=4.0*cm,
                    preserveAspectRatio=True, mask='auto')
        c.setFont("Helvetica", 9)
        c.drawString(xi, yi - 0.3*cm,
                     f"Upload {i+1} — {display_type(f.get('file_type'))}")

    draw_footer(c, page)
    c.showPage(); page += 1

    draw_header(c, title)
    y = H - 3.4*cm

    y = section_title(c, "Demographic Details", MARGIN, y)
    demo = case.get("demographic_details", {}) or {}
    ci = demo.get("contact_information", {}) or {}
    pd = demo.get("personal_details", {}) or {}
    label_value(c, MARGIN, y, "Aadhaar Photo Match", demo.get("aadhaar_photo_match")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Email ID", ci.get("email_id")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Mobile Number", ci.get("mobile_number")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Gender", pd.get("gender")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Education", pd.get("education")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Family Members", pd.get("number_of_family_members")); y -= 1.0*cm

    y = section_title(c, "Business Details", MARGIN, y)
    bd = case.get("business_details", {}) or {}
    ei  = bd.get("enterprise_information", {}) or {}
    bld = bd.get("business_location_details", {}) or {}
    ba  = bd.get("business_address", {}) or {}
    bact= bd.get("business_activity", {}) or {}
    binfo=bd.get("business_info", {}) or {}
    label_value(c, MARGIN, y, "Enterprise Name", ei.get("enterprise_name")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Organization Type", ei.get("type_of_organization")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Business Location Type", bld.get("business_location")); y -= 0.5*cm
    label_value(c, MARGIN, y, "Address", ba.get("address")); y -= 0.5*cm
    label_value(c, MARGIN, y,
                "City/District/State/Pincode",
                f"{ba.get('city','')}, {ba.get('district','')}, {ba.get('state','')} - {ba.get('pincode','')}")
    y -= 1.0*cm

    y = section_title(c, "Financial Details", MARGIN, y)
    fin  = case.get("financial_details", {}) or {}
    bfi  = fin.get("business_financial_information", {}) or {}
    ffi  = fin.get("family_financial_information", {}) or {}
    def fmt_inr(v):
        try:
            return f"{int(v):,} INR"
        except:
            return "N/A"

    label_value(c, MARGIN, y, "Monthly Income", fmt_inr(bfi.get("monthly_income_from_business"))); y -= 0.5*cm
    label_value(c, MARGIN, y, "Monthly Expense", fmt_inr(bfi.get("monthly_expense_of_business"))); y -= 0.5*cm
    label_value(c, MARGIN, y, "Family Income", fmt_inr(ffi.get("monthly_family_income"))); y -= 0.5*cm

    draw_footer(c, page)
    c.showPage(); page += 1

    for idx, f in enumerate(files, start=1):
        draw_header(c, title)
        y = H - 3.4*cm

        c.setFont("Helvetica-Bold", 14)
        c.drawString(MARGIN, y, f"Upload {idx} — {display_type(f.get('file_type'))}")
        y -= 0.8*cm

                # ✅ INTEGRITY VERIFICATION DASHBOARD IMAGE - AT THE TOP BELOW TITLE
        try:
            # Load and display the integrity verification image
            integrity_image_path = "integrity_dashboard.png"  # Make sure this file is in your Docker container
            if os.path.exists(integrity_image_path):
                # Position the image right below the title
                integrity_y = y
                
                # Display the image - centered and sized appropriately
                image_width = 16*cm  # Adjust width as needed
                image_height = 2.5*cm  # Adjust height proportionally
                
                # Calculate X position to center the image
                center_x = (W - image_width) / 2
                
                c.drawImage(
                    integrity_image_path,
                    center_x,  # Centered horizontally
                    integrity_y - image_height,  # Position below title
                    width=image_width,
                    height=image_height,
                    preserveAspectRatio=True,
                    mask='auto'
                )
                
                # Adjust the Y position for the following content (image + some spacing)
                y = integrity_y - image_height - 0.8*cm
            else:
                print(f"Warning: Integrity image not found at {integrity_image_path}")
                y -= 0.5*cm
        except Exception as e:
            print(f"Warning: Could not load integrity image: {e}")
            y -= 0.5*cm
            
        gl = f.get("geo_location", {}) or {}
        lat = gl.get("lat")
        lng = gl.get("lng")

        img = make_thumb(f["_bin"])
        c.drawImage(ImageReader(img), MARGIN, y-9.8*cm,
                    width=8.6*cm, height=9.0*cm, preserveAspectRatio=True, mask='auto')

        mimg = static_map_with_path(lat, lng, case_lat, case_lng, zoom=17)
        if mimg:
            c.drawImage(ImageReader(mimg), 11.0*cm, y-9.8*cm,
                        width=8.0*cm, height=9.0*cm,
                        preserveAspectRatio=True, mask='auto')

        yb = y - 10.4*cm
        rev = reverse_geocode(lat, lng)
        yb = draw_wrapped_label_value(c, MARGIN, yb, "Location Address", rev, max_width_cm=18.0)

        dist_km = haversine_km(case_lat, case_lng, lat, lng) if (lat is not None) else None
        dist_m  = round(dist_km * 1000) if dist_km else None
        c.setFont("Helvetica", 10)
        c.drawString(MARGIN, yb,
            f"Distance from case location: {dist_m} meters" if dist_m else "Distance from case location: N/A")
        yb -= 0.7*cm

        up_at = fmt_ist(f.get("uploaded_at"))
        c.setFont("Helvetica", 10)
        c.drawString(MARGIN, yb, f"Uploaded At: {up_at}")
        yb -= 1.0*cm

        c.setFont("Helvetica-Bold", 12)
        c.drawString(MARGIN, yb, "Metadata & Sensor Information")
        yb -= 0.6*cm
        c.setFont("Helvetica", 10)
        c.drawString(MARGIN, yb, f"Server Time: {up_at}")
        yb -= 0.6*cm
        c.drawString(MARGIN, yb, f"Device Time: {up_at}")
        yb -= 0.6*cm
        c.drawString(MARGIN, yb,
            f"Accuracy: {dist_m} meters" if dist_m else "Distance from case location: N/A")
        yb -= 0.6*cm
        if lat:
            c.drawString(MARGIN, yb, f"Geo: Lat {lat:.6f}, Lng {lng:.6f}")
        else:
            c.drawString(MARGIN, yb, "Geo: N/A")
        yb -= 0.6*cm

        draw_footer(c, page)
        c.showPage(); page += 1

    draw_header(c, title)
    y = H - 3.4*cm
    c.setFont("Helvetica-Bold", 14)
    c.drawString(MARGIN, y, "Inspection Timeline")
    y -= 0.8*cm

    c.setFillColorRGB(0.88, 0.97, 0.90)
    c.roundRect(MARGIN, y-0.8*cm, W-2*MARGIN, 1.1*cm, 0.3*cm, fill=1, stroke=0)
    c.setFillColorRGB(0.12, 0.45, 0.25)
    c.setFont("Helvetica-Bold", 12)
    c.drawString(MARGIN + 0.6*cm, y-0.4*cm, "●  Ready for Review")
    c.setFillColor(colors.black)
    y -= 2.0*cm

    events = []
    def evt(text, ts, sub=None):
        if ts:
            events.append((text, ts, sub))

    evt("Case Created", case.get("created_at"))
    evt("Assigned to Agent", case.get("assigned_at"))
    evt("Accepted by Agent", case.get("accepted_at"))

    for f in files:
        sub = reverse_geocode(f.get("geo_location", {}).get("lat"), f.get("geo_location", {}).get("lng"))
        evt(f"Upload — {display_type(f.get('file_type'))}",
            f.get("uploaded_at"),
            sub)

    evt("Case Updated", case.get("updated_at"))
    evt("Case Completed", case.get("completed_at"))

    def parse(s):
        try:
            return datetime.fromisoformat(s.replace("Z","+00:00"))
        except:
            return datetime.min
    events.sort(key=lambda x: parse(x[1]))

    for text, ts, sub in events:
        c.setFont("Helvetica", 10)
        c.drawString(MARGIN, y, f"• {fmt_ist(ts)}")
        y -= 0.45*cm

        c.setFont("Helvetica-Bold", 11)
        c.drawString(MARGIN + 0.6*cm, y, text)
        y -= 0.45*cm

        if sub:
            c.setFont("Helvetica", 9)
            c.setFillColor(colors.grey)
            y = draw_wrapped_label_value(c, MARGIN + 0.6*cm, y, "Uploaded from", sub, max_width_cm=17.0)
            c.setFillColor(colors.black)
            y += 0.3*cm

        y -= 0.3*cm

        if y < 3*cm:
            draw_footer(c, page)
            c.showPage(); page += 1
            draw_header(c, title)
            y = H - 3.4*cm
            c.setFont("Helvetica-Bold", 14)
            c.drawString(MARGIN, y, "Inspection Timeline (Cont.)")
            y -= 1.0*cm

    draw_footer(c, page)
    c.showPage()
    c.save()

# FastAPI endpoints
class CaseRequest(BaseModel):
    case_id: str

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

@app.post("/generate-report")
async def generate_report(case_request: CaseRequest):
    try:
        case_id = case_request.case_id
        print(f"🔐 Logging in for case {case_id}...")
        admin_login()

        print("📂 Fetching case…")
        case = get_case(case_id)

        print("📁 Fetching files…")
        files = get_files(case_id)

        print("📸 Downloading binaries…")
        for f in files:
            f["_bin"] = download_file(f["file_id"])

        print("🧭 Generating report…")
        output_filename = f"Report_{case_id}.pdf"
        output_path = f"/app/output/{output_filename}"
        
        build_pdf(case, files, output_path)
        
        with open(output_path, "rb") as f:
            pdf_content = f.read()
        
        # Clean up
        if os.path.exists(output_path):
            os.remove(output_path)
            
        return Response(
            content=pdf_content,
            media_type="application/pdf",
            headers={"Content-Disposition": f"attachment; filename={output_filename}"}
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating PDF: {str(e)}")

@app.get("/")
async def root():
    return {"message": "Case Report Generator API", "status": "running"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)