"""
=============================================================================
  POLE ENGINEERING PROJECT AUTOMATION SCRIPT
  Automatisation des projets d'ingénierie de poteaux utilitaires
=============================================================================
  Workflow:
    1. Create project folder structure
    2. Read source data from "départ" folder
    3. Rename & move folders to RSA-UDS with business logic (sd/d)
    4. Identify pole/anchor images via Gemini AI
    5. Convert images & Excel files to PDF
    6. Move all PDFs to compilation folder
=============================================================================
  REQUIREMENTS:
    pip install pillow openpyxl reportlab google-generativeai
=============================================================================
"""
 
import os
import re
import sys
import shutil
import base64
import argparse
import logging
from pathlib import Path
 
# ── Third-party ──────────────────────────────────────────────────────────────
try:
    from PIL import Image
except ImportError:
    print("❌  Pillow missing. Run: pip install pillow")
    sys.exit(1)
 
try:
    import openpyxl
    from openpyxl import load_workbook
except ImportError:
    print("❌  openpyxl missing. Run: pip install openpyxl")
    sys.exit(1)
 
try:
    from reportlab.lib.pagesizes import letter, landscape
    from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
    from reportlab.lib.styles import getSampleStyleSheet
    from reportlab.lib import colors
    from reportlab.pdfgen import canvas
except ImportError:
    print("❌  reportlab missing. Run: pip install reportlab")
    sys.exit(1)
 
try:
    import google.generativeai as genai
    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False
    print("⚠️   google-generativeai not installed. Gemini AI disabled.")
    print("     Run: pip install google-generativeai")
 
# ── Logging ───────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)
 
# ═════════════════════════════════════════════════════════════════════════════
#  CONFIGURATION  –  edit these paths to match your environment
# ═════════════════════════════════════════════════════════════════════════════
DEFAULT_BASE_DIR   = Path.home() / "PycharmProjects" / "structure" / "projet"  # Where projects are created
DEFAULT_DEPART_DIR = Path.home() / "PycharmProjects" / "structure" / "départ"   # Source folder with pole sub-folders
GEMINI_API_KEY     = os.environ.get("AIzaSyDq2LmX_fwKGAwxGAtmBfX940vT2wDQzBU")  # Set via env var
GEMINI_MODEL       = "models/gemini-2.5-flash"
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  STEP 1 – PROJECT FOLDER CREATION
# ═════════════════════════════════════════════════════════════════════════════
 
SUB_FOLDERS = ["compilation", "ingénierie", "RSA-UDS", "plan"]
 
def create_project_structure(project_name: str, base_dir: Path) -> Path:
    """
    Create:
        <base_dir>/<project_name>/
            compilation/
            ingénierie/
            RSA-UDS/
            plan/
    Returns the project root path.
    """
    project_root = base_dir / project_name
    project_root.mkdir(parents=True, exist_ok=True)
 
    for folder in SUB_FOLDERS:
        (project_root / folder).mkdir(exist_ok=True)
 
    log.info(f"✅  Project structure created: {project_root}")
    return project_root
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  STEP 2 – PARSE SOURCE POLE FOLDERS
# ═════════════════════════════════════════════════════════════════════════════
 
# Pattern:  <seq><d|sd><anchors>-<pole_ref> <address>
#   e.g.    1d1-PBO345 rue saint-eugene
#           2sd2-PBO678 boul. des érables
FOLDER_PATTERN = re.compile(
    r"^(?P<seq>\d+)"
    r"(?P<dero>sd|d)"
    r"(?P<anchors>\d+)"
    r"-"
    r"(?P<pole_ref>[^\s]+)"
    r"(?:\s+(?P<address>.+))?$",
    re.IGNORECASE,
)
 
def parse_pole_folder_name(folder_name: str) -> dict | None:
    """
    Parse a source folder name and return a dict with extracted components,
    or None if the name doesn't match the expected pattern.
    """
    m = FOLDER_PATTERN.match(folder_name)
    if not m:
        return None
    return {
        "seq"      : int(m.group("seq")),
        "dero"     : m.group("dero").lower(),   # 'sd' or 'd'
        "anchors"  : int(m.group("anchors")),
        "pole_ref" : m.group("pole_ref"),        # e.g. PBO345
        "address"  : (m.group("address") or "").strip(),
        "original" : folder_name,
    }
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  STEP 3 – RENAMING & MOVING LOGIC (sd vs d)
# ═════════════════════════════════════════════════════════════════════════════
 
def compute_new_folder_name(parsed: dict, current_seq: int) -> str:
    """
    Apply business naming rules:
      sd  →  10-PBO345              (prefix = current_seq, no address)
      d   →  11-PBO345 rue saint-x  (prefix = current_seq, with address)
 
    NOTE: The leading '1' before 0/1 comes from your spec (10- / 11-).
          If you want the sequence itself used as prefix, adjust here.
    """
    if parsed["dero"] == "sd":
        return f"{current_seq}0-{parsed['pole_ref']}"
    else:
        addr = f" {parsed['address']}" if parsed["address"] else ""
        return f"{current_seq}1-{parsed['pole_ref']}{addr}"
 
 
def rename_excel_files(folder_path: Path, new_folder_name: str):
    """Rename any .xlsx / .xls inside folder to match new_folder_name."""
    for f in folder_path.iterdir():
        if f.suffix.lower() in (".xlsx", ".xls"):
            new_name = new_folder_name + f.suffix
            f.rename(folder_path / new_name)
            log.info(f"   📄  Excel renamed → {new_name}")
 
 
def process_depart_to_rsa_uds(depart_dir: Path, rsa_uds_dir: Path) -> list[Path]:
    """
    Read all sub-folders in départ, sort by sequence number,
    apply renaming logic, move to RSA-UDS.
    Returns list of destination folder paths (for further processing).
    """
    entries = []
    for item in depart_dir.iterdir():
        if item.is_dir():
            parsed = parse_pole_folder_name(item.name)
            if parsed:
                entries.append((parsed, item))
            else:
                log.warning(f"⚠️   Unrecognised folder name (skipped): {item.name}")
 
    # Sort by original sequence
    entries.sort(key=lambda x: x[0]["seq"])
 
    current_seq = 1
    moved_folders = []
 
    for parsed, src_path in entries:
        new_name = compute_new_folder_name(parsed, current_seq)
        dst_path = rsa_uds_dir / new_name
 
        shutil.copytree(src_path, dst_path)
        log.info(f"📁  {src_path.name!r:40s} → {new_name!r}")
 
        # Rename Excel files inside the destination folder
        rename_excel_files(dst_path, new_name)
 
        moved_folders.append(dst_path)
 
        # Advance counter
        if parsed["dero"] == "sd":
            current_seq += 1
        else:
            current_seq += 2   # Extra jump for CDC document
 
    return moved_folders
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  STEP 4 – GEMINI AI IMAGE IDENTIFICATION
# ═════════════════════════════════════════════════════════════════════════════
 
def _encode_image(image_path: Path) -> str:
    """Return base64-encoded image content."""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")
 
 
def identify_images_with_gemini(folder_path: Path) -> dict:
    """
    Send all JPG images in folder to Gemini and ask it to label each as
    'pole' or 'anchor'.
 
    Returns:
        {
            "pole":   ["IMG_001.jpg", ...],
            "anchor": ["IMG_002.jpg", ...],
            "unknown": [...]
        }
    """
    jpg_files = sorted(folder_path.glob("*.jpg")) + sorted(folder_path.glob("*.JPG")) + \
                sorted(folder_path.glob("*.jpeg"))
 
    if not jpg_files:
        log.warning(f"   No JPG images found in {folder_path.name}")
        return {"pole": [], "anchor": [], "unknown": []}
 
    if not GEMINI_AVAILABLE or not GEMINI_API_KEY:
        log.warning("   Gemini unavailable – using filename heuristic fallback.")
        return _fallback_image_classification(jpg_files)
 
    genai.configure(api_key=GEMINI_API_KEY)
    model = genai.GenerativeModel(GEMINI_MODEL)
 
    # Build the content list
    file_names = [f.name for f in jpg_files]
    prompt = (
        "You are analysing images for a utility pole engineering project. "
        "For each image provided, classify it as either 'pole' or 'anchor'. "
        "A pole image shows a utility/telephone pole. "
        "An anchor image shows a ground anchor or guy-wire anchor attachment. "
        f"The filenames are: {file_names}. "
        "Respond ONLY with a JSON object where each key is a filename and "
        "the value is either 'pole' or 'anchor'. Example: "
        '{"IMG_001.jpg": "pole", "IMG_002.jpg": "anchor"}'
    )
 
    content = [prompt]
    for img_path in jpg_files:
        content.append({
            "mime_type": "image/jpeg",
            "data": _encode_image(img_path),
        })
 
    try:
        response = model.generate_content(content)
        raw = response.text.strip()
        # Strip markdown code fences if present
        raw = re.sub(r"^```json\s*|```$", "", raw, flags=re.MULTILINE).strip()
        import json
        classification = json.loads(raw)
    except Exception as e:
        log.error(f"   Gemini error: {e} – falling back to heuristic.")
        return _fallback_image_classification(jpg_files)
 
    result = {"pole": [], "anchor": [], "unknown": []}
    for fname, label in classification.items():
        label = label.lower()
        if label in result:
            result[label].append(fname)
        else:
            result["unknown"].append(fname)
 
    log.info(f"   🤖  Gemini classified: pole={result['pole']}, anchor={result['anchor']}")
    return result
 
 
def _fallback_image_classification(jpg_files: list[Path]) -> dict:
    """
    Simple heuristic fallback when Gemini is unavailable.
    Files with 'pole' or 'p' in name → pole; 'anchor' or 'a' → anchor.
    Remaining: first half = pole, second half = anchor.
    """
    result = {"pole": [], "anchor": [], "unknown": []}
    unclassified = []
 
    for f in jpg_files:
        lower = f.name.lower()
        if "pole" in lower or lower.startswith("p"):
            result["pole"].append(f.name)
        elif "anchor" in lower or lower.startswith("a"):
            result["anchor"].append(f.name)
        else:
            unclassified.append(f.name)
 
    mid = len(unclassified) // 2
    result["pole"].extend(unclassified[:mid])
    result["anchor"].extend(unclassified[mid:])
 
    return result
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  STEP 5 – FILE CONVERSION & FINAL SORTING
# ═════════════════════════════════════════════════════════════════════════════
 
def images_to_pdf(image_paths: list[Path], output_pdf: Path):
    """Combine one or more JPG images into a single PDF page-per-image."""
    if not image_paths:
        return
 
    pil_images = []
    for p in image_paths:
        img = Image.open(p).convert("RGB")
        pil_images.append(img)
 
    if pil_images:
        first = pil_images[0]
        rest  = pil_images[1:]
        first.save(output_pdf, save_all=True, append_images=rest)
        log.info(f"   🖼→📄  {[p.name for p in image_paths]} → {output_pdf.name}")
 
 
def excel_to_pdf(excel_path: Path, output_pdf: Path):
    """
    Convert an Excel file to a PDF using ReportLab.
    Reads all sheets and renders them as tables.
    """
    wb = load_workbook(excel_path, data_only=True)
    doc = SimpleDocTemplate(
        str(output_pdf),
        pagesize=landscape(letter),
        rightMargin=20, leftMargin=20,
        topMargin=30, bottomMargin=20,
    )
    styles = getSampleStyleSheet()
    story  = []
 
    for sheet_name in wb.sheetnames:
        ws = wb[sheet_name]
        story.append(Paragraph(f"Sheet: {sheet_name}", styles["Heading2"]))
        story.append(Spacer(1, 6))
 
        data = []
        for row in ws.iter_rows(values_only=True):
            data.append([str(cell) if cell is not None else "" for cell in row])
 
        if not data:
            continue
 
        col_count = max(len(r) for r in data)
        col_width  = (landscape(letter)[0] - 40) / max(col_count, 1)
 
        table = Table(data, colWidths=[col_width] * col_count, repeatRows=1)
        table.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#2C4770")),
            ("TEXTCOLOR",  (0, 0), (-1, 0), colors.white),
            ("FONTSIZE",   (0, 0), (-1, -1), 7),
            ("GRID",       (0, 0), (-1, -1), 0.25, colors.grey),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.white, colors.HexColor("#F0F4FA")]),
            ("VALIGN",     (0, 0), (-1, -1), "MIDDLE"),
        ]))
        story.append(table)
        story.append(Spacer(1, 20))
 
    doc.build(story)
    log.info(f"   📊→📄  {excel_path.name} → {output_pdf.name}")
 
 
def process_pole_folder(
    pole_folder: Path,
    compilation_dir: Path,
    pdf_seq: int,
) -> int:
    """
    For a single pole folder (already in RSA-UDS):
      1. Identify images via Gemini
      2. Convert pole images → PDF  (pdf_seq)
      3. Convert anchor images → PDF (pdf_seq + 1)
      4. Convert Excel → PDF         (pdf_seq + 2)
      5. Move all PDFs to compilation
 
    Returns the next available pdf_seq.
    """
    folder_label = pole_folder.name   # e.g. "10-PBO345"
 
    # Strip leading numeric prefix for base label  e.g. "PBO345"
    base_label = re.sub(r"^\d+-", "", folder_label)
 
    # ── Identify images ──────────────────────────────────────────────────────
    classification = identify_images_with_gemini(pole_folder)
 
    # ── Pole PDF ─────────────────────────────────────────────────────────────
    pole_images = [pole_folder / n for n in classification["pole"] if (pole_folder / n).exists()]
    if pole_images:
        pole_pdf_name = f"{pdf_seq}-{base_label}.pdf"
        images_to_pdf(pole_images, compilation_dir / pole_pdf_name)
        pdf_seq += 1
 
    # ── Anchor PDF ───────────────────────────────────────────────────────────
    anchor_images = [pole_folder / n for n in classification["anchor"] if (pole_folder / n).exists()]
    if anchor_images:
        anchor_pdf_name = f"{pdf_seq}-{base_label}.pdf"
        images_to_pdf(anchor_images, compilation_dir / anchor_pdf_name)
        pdf_seq += 1
 
    # ── Excel → PDF ──────────────────────────────────────────────────────────
    for xlsx_file in list(pole_folder.glob("*.xlsx")) + list(pole_folder.glob("*.xls")):
        excel_pdf_name = f"{pdf_seq}-{base_label}.pdf"
        excel_to_pdf(xlsx_file, compilation_dir / excel_pdf_name)
        pdf_seq += 1
 
    return pdf_seq
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  MAIN ORCHESTRATOR
# ═════════════════════════════════════════════════════════════════════════════
 
def run(project_name: str, base_dir: Path, depart_dir: Path):
    log.info("=" * 65)
    log.info(f"  PROJECT : {project_name}")
    log.info(f"  BASE DIR: {base_dir}")
    log.info(f"  DÉPART  : {depart_dir}")
    log.info("=" * 65)
 
    # 1 – Create structure
    project_root   = create_project_structure(project_name, base_dir)
    rsa_uds_dir    = project_root / "RSA-UDS"
    compilation_dir = project_root / "compilation"
 
    # 2 & 3 – Move and rename pole folders
    if not depart_dir.exists():
        log.error(f"Source départ folder not found: {depart_dir}")
        sys.exit(1)
 
    moved_folders = process_depart_to_rsa_uds(depart_dir, rsa_uds_dir)
    log.info(f"📂  {len(moved_folders)} pole folder(s) moved to RSA-UDS.")
 
    # 4 & 5 – Process each pole folder
    pdf_seq = 1
    for folder in sorted(moved_folders):
        log.info(f"\n🔧  Processing: {folder.name}")
        pdf_seq = process_pole_folder(folder, compilation_dir, pdf_seq)
 
    log.info("\n" + "=" * 65)
    log.info(f"✅  DONE.  All PDFs are in: {compilation_dir}")
    log.info("=" * 65)
 
 
# ═════════════════════════════════════════════════════════════════════════════
#  CLI ENTRY POINT
# ═════════════════════════════════════════════════════════════════════════════
 
def main():
    parser = argparse.ArgumentParser(
        description="Utility pole engineering project automation script",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python pole_automation.py E65-BK7-26-132679
  python pole_automation.py E65-BK7-26-132679 --base-dir /Users/me/Projects
  python pole_automation.py E65-BK7-26-132679 --depart-dir /Users/me/départ
  python pole_automation.py E65-BK7-26-132679 --gemini-key AIza...
        """,
    )
    parser.add_argument("project_name", help="Project name (e.g. E65-BK7-26-132679)")
    parser.add_argument("--base-dir",   type=Path, default=DEFAULT_BASE_DIR,
                        help=f"Root directory for projects (default: {DEFAULT_BASE_DIR})")
    parser.add_argument("--depart-dir", type=Path, default=DEFAULT_DEPART_DIR,
                        help=f"Source 'départ' folder (default: {DEFAULT_DEPART_DIR})")
    parser.add_argument("--gemini-key", type=str, default=GEMINI_API_KEY,
                        help="Gemini API key (or set GEMINI_API_KEY env var)")
 
    args = parser.parse_args()
 
    # Inject Gemini key if provided via CLI
    if args.gemini_key:
        global GEMINI_API_KEY
        GEMINI_API_KEY = args.gemini_key
 
    run(args.project_name, args.base_dir, args.depart_dir)
 
 
if __name__ == "__main__":
    main()