import os
import re
import sys
import shutil
import base64
import argparse
import logging
import json
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
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

# -- Logging --
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s", datefmt="%H:%M:%S")
log = logging.getLogger(__name__)

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", "AIzaSyDq2LmX_fwKGAwxGAtmBfX940vT2wDQzBU")  # Set via env var
GEMINI_MODEL       = "models/gemini-2.5-flash"

FOLDER_PATTERN = re.compile(r"^(?P<seq>\d+)(?P<dero>sd|d)(?P<anchors>\d+)-(?P<pole_ref>[^\s]+)(?:\s+(?P<address>.+))?$", re.IGNORECASE)

def create_project_structure(project_name: str, base_dir: Path) -> Path:
    project_root = base_dir / project_name
    for folder in ["compilation", "ingénierie", "RSA-UDS", "plan"]:
        (project_root / folder).mkdir(parents=True, exist_ok=True)
    return project_root

def parse_pole_folder_name(folder_name: str) -> dict | None:
    m = FOLDER_PATTERN.match(folder_name)
    if not m: return None
    return {
        "seq": int(m.group("seq")),
        "dero": m.group("dero").lower(),
        "anchors": int(m.group("anchors")),
        "pole_ref": m.group("pole_ref"),
        "address": (m.group("address") or "").strip(),
        "original": folder_name,
    }

def compute_new_folder_name(parsed: dict, current_seq: int) -> str:
    # Always starts with the sequence provided (1, 3, 5...)
    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):
    for f in folder_path.iterdir():
        if f.suffix.lower() in (".xlsx", ".xls"):
            # According to your request, UDS (Excel) should be named 10-
            # We will handle the specific 10- prefix during the PDF conversion in process_pole_folder
            pass

def identify_images_with_gemini(folder_path: Path) -> dict:
    jpg_files = sorted(list(folder_path.glob("*.jpg")) + list(folder_path.glob("*.JPG")) + list(folder_path.glob("*.jpeg")))
    if not jpg_files: return {"pole": [], "anchor": []}
    
    if not GEMINI_AVAILABLE:
        return _fallback_image_classification(jpg_files)

    genai.configure(api_key=GEMINI_API_KEY)
    model = genai.GenerativeModel(GEMINI_MODEL)
    
    file_names = [f.name for f in jpg_files]
    prompt = f"Classify these images as 'pole' or 'anchor'. Files: {file_names}. Respond ONLY in JSON."
    
    content = [prompt]
    for img_path in jpg_files:
        with open(img_path, "rb") as f:
            content.append({"mime_type": "image/jpeg", "data": base64.b64encode(f.read()).decode("utf-8")})

    try:
        response = model.generate_content(content)
        raw = re.sub(r"^```json\s*|```$", "", response.text, flags=re.MULTILINE).strip()
        classification = json.loads(raw)
        result = {"pole": [], "anchor": []}
        for fname, label in classification.items():
            if label.lower() in result: result[label.lower()].append(fname)
        return result
    except:
        return _fallback_image_classification(jpg_files)

def _fallback_image_classification(jpg_files):
    mid = len(jpg_files) // 2
    return {"pole": [f.name for f in jpg_files[:mid]], "anchor": [f.name for f in jpg_files[mid:]]}

#--------------------------------------------
# Transforme les photos en pdf
#--------------------------------------------

def images_to_pdf(image_paths: list[Path], output_pdf: Path):
    pil_images = [Image.open(p).convert("RGB") for p in image_paths if p.exists()]
    if pil_images:
        pil_images[0].save(output_pdf, save_all=True, append_images=pil_images[1:])

def excel_to_pdf(excel_path: Path, output_pdf: Path):
    wb = load_workbook(excel_path, data_only=True)
    doc = SimpleDocTemplate(str(output_pdf), pagesize=landscape(letter))
    styles = getSampleStyleSheet()
    story = []
    for sheet_name in wb.sheetnames:
        ws = wb[sheet_name]
        data = [[str(cell) if cell is not None else "" for cell in row] for row in ws.iter_rows(values_only=True)]
        if not data: continue
        t = Table(data, repeatRows=1)
        t.setStyle(TableStyle([('FONTSIZE', (0,0), (-1,-1), 7), ('GRID', (0,0), (-1,-1), 0.5, colors.grey)]))
        story.append(t)
    doc.build(story)

def create_dummy_cdc(output_path: Path, label: str):
    """Creates a placeholder PDF for the CDC requirement."""
    c = canvas.Canvas(str(output_path), pagesize=letter)
    c.drawString(100, 750, f"CDC Document for {label}")
    c.save()

#---------------------------
#This function is the "Orchestrator" for each individual pole. 
#It decides the order of your documents and ensures the numbering matches your specific requirements.
#-----------------------------

def process_pole_folder(pole_folder: Path, compilation_dir: Path, folder_parsed: dict, current_pdf_idx: int) -> int:
    base_label = folder_parsed['pole_ref']
    address = f" {folder_parsed['address']}" if folder_parsed['address'] else ""
    full_label = f"{base_label}{address}" # Includes address if present (e.g. PHA627-3 Route 132)
    
    # 1. Identify Images
    classification = identify_images_with_gemini(pole_folder)
    
    # -- A. CDC (Calcul de Charge) - Only if 'd' --
    if folder_parsed["dero"] == "d":
        cdc_name = f"{current_pdf_idx}-CDC-{full_label}.pdf"
        # Logic to create/move CDC goes here
        log.info(f"   📑 CDC assigned: {cdc_name}")
        current_pdf_idx += 1

    # -- B. UDS (Excel) -> Example: 10-RSA PHA627-1.pdf --
    for xlsx_file in list(pole_folder.glob("*.xlsx")) + list(pole_folder.glob("*.xls")):
        excel_pdf_name = f"{current_pdf_idx}-RSA {full_label}.pdf"
        excel_to_pdf(xlsx_file, compilation_dir / excel_pdf_name)
        log.info(f"   📊 UDS Excel → {excel_pdf_name}")
        current_pdf_idx += 1

    # -- C. Poteau (Pole Images) -> Example: 11-PHOTO PHA627-1.pdf --
    pole_images = [pole_folder / n for n in classification["pole"]]
    if pole_images:
        pole_pdf_name = f"{current_pdf_idx}-PHOTO {full_label}.pdf"
        images_to_pdf(pole_images, compilation_dir / pole_pdf_name)
        log.info(f"   🖼 Poteau → {pole_pdf_name}")
        current_pdf_idx += 1

    # -- D. Ancre (Anchor Images) -> Example: 12-PHOTO ANCRE PHA627-1.pdf --
    anchor_images = [pole_folder / n for n in classification["anchor"]]
    if anchor_images:
        anchor_pdf_name = f"{current_pdf_idx}-PHOTO ANCRE {full_label}.pdf"
        images_to_pdf(anchor_images, compilation_dir / anchor_pdf_name)
        log.info(f"   ⚓ Ancre → {anchor_pdf_name}")
        current_pdf_idx += 1

    # We return the new index so the NEXT pole knows where to start
    return current_pdf_idx

def run(project_name: str, base_dir: Path, depart_dir: Path):
    # 1. Create the project structure
    project_root = create_project_structure(project_name, base_dir)
    rsa_uds_dir = project_root / "RSA-UDS"
    compilation_dir = project_root / "compilation"

    # --- ADDED: INITIALIZE ENTRIES ---
    entries = []
    if not depart_dir.exists():
        log.error(f"❌ Source folder not found: {depart_dir}")
        return

    # Loop through the 'départ' folder to find pole sub-folders
    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"⚠️  Skipping folder (unrecognized format): {item.name}")
    # ---------------------------------

    # 2. Sort folders to ensure 1sd1 comes before 2d2
    entries.sort(key=lambda x: x[0]["seq"])
    
    # GLOBAL PDF COUNTER (Starts at 10)
    global_pdf_count = 10 

    for parsed, src_path in entries:
        # 3. Create the RSA-UDS folder name
        new_folder_name = compute_new_folder_name(parsed, parsed["seq"]) 
        dst_path = rsa_uds_dir / new_folder_name
        
        # Avoid error if folder already exists
        if dst_path.exists():
            shutil.rmtree(dst_path)
            
        shutil.copytree(src_path, dst_path)

        # 4. Process documents and update the global counter
        log.info(f"\n🔧 Processing folder: {new_folder_name}")
        global_pdf_count = process_pole_folder(dst_path, compilation_dir, parsed, global_pdf_count)
        
    log.info(f"\n✅ All done! PDFs generated in: {compilation_dir}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("project_name")
    parser.add_argument("--base-dir", type=Path, default=DEFAULT_BASE_DIR)
    parser.add_argument("--depart-dir", type=Path, default=DEFAULT_DEPART_DIR)
    args = parser.parse_args()
    run(args.project_name, args.base_dir, args.depart_dir)