#!/usr/bin/env python3
"""
GAVEL — Privilege Log Generator
=================================
Scans a folder of documents and generates a privilege log spreadsheet.
Uses Claude AI to write a one-line description of each document.

SETUP (run once):
    Mac:     /usr/local/bin/pip3 install anthropic pypdf openpyxl
    Windows: pip install anthropic pypdf openpyxl

USAGE:
    Mac:     python3 privilege_log.py
    Windows: python privilege_log.py

    The script will ask you for your API key, folder, and options.
    No editing required.

OUTPUT:
    An Excel spreadsheet (.xlsx) with these columns:
    - Document Number
    - Filename
    - File Type
    - Date Modified
    - File Size
    - Description (generated by Claude AI)
    - Privilege Claimed
    - Notes
"""

import os
import platform
from datetime import datetime

# ---------------------------------------------------------------
# OPTIONAL — customize the privilege types listed in the dropdown
# ---------------------------------------------------------------
PRIVILEGE_TYPES = [
    "Attorney-Client Privilege",
    "Work Product Doctrine",
    "Attorney-Client Privilege and Work Product Doctrine",
    "Confidential",
    "Not Privileged",
]
# ---------------------------------------------------------------


def check_libraries():
    """Check required libraries are installed."""
    missing = []
    for lib in ['anthropic', 'pypdf', 'openpyxl']:
        try:
            __import__(lib)
        except ImportError:
            missing.append(lib)

    if missing:
        print("\nERROR: Missing required libraries: " + ", ".join(missing))
        print("Install them by running:")
        if platform.system() == 'Windows':
            print(f"    pip install {' '.join(missing)}")
        else:
            print(f"    /usr/local/bin/pip3 install {' '.join(missing)}")
        return False
    return True


def get_input():
    """Ask the user for API key, folder, and options."""
    print("\n" + "-" * 60)

    # Get API key
    api_key = input("Enter your Anthropic API key (starts with sk-ant-): ").strip()
    if not api_key.startswith("sk-ant-"):
        print("  Warning: API key doesn't look right — make sure you copied the full key.")

    # Get folder
    while True:
        folder = input("Enter the full path to the folder containing your documents: ").strip()
        folder = folder.strip('"').strip("'")
        folder = os.path.normpath(os.path.expanduser(folder))
        if os.path.exists(folder):
            break
        print(f"  Folder not found: {folder}")

    # Get output filename
    default_name = f"privilege_log_{datetime.now().strftime('%Y-%m-%d')}.xlsx"
    name_input = input(f"Output filename (press Enter for '{default_name}'): ").strip()
    output_name = name_input if name_input else default_name
    if not output_name.lower().endswith('.xlsx'):
        output_name += '.xlsx'

    # Get case name
    case_input = input("Case name (optional, press Enter to skip): ").strip()

    return api_key, folder, output_name, case_input


def get_file_info(filepath):
    """Get metadata for a file."""
    stat = os.stat(filepath)
    size_bytes = stat.st_size
    modified = datetime.fromtimestamp(stat.st_mtime)

    # Format size
    if size_bytes < 1024:
        size_str = f"{size_bytes} B"
    elif size_bytes < 1024 * 1024:
        size_str = f"{size_bytes / 1024:.1f} KB"
    else:
        size_str = f"{size_bytes / (1024*1024):.1f} MB"

    return {
        'size': size_str,
        'modified': modified.strftime("%B %d, %Y"),
        'extension': os.path.splitext(filepath)[1].upper().replace('.', '')
    }


def extract_text(filepath):
    """Extract text from a file for Claude to analyze."""
    ext = os.path.splitext(filepath)[1].lower()
    text = ""

    try:
        if ext == '.pdf':
            from pypdf import PdfReader
            reader = PdfReader(filepath)
            # Extract first 3 pages max to keep API costs low
            for i, page in enumerate(reader.pages):
                if i >= 3:
                    break
                text += page.extract_text() or ""

        elif ext in ['.txt', '.md', '.csv']:
            with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
                text = f.read(3000)  # First 3000 chars

        elif ext in ['.docx']:
            try:
                import zipfile
                import xml.etree.ElementTree as ET
                with zipfile.ZipFile(filepath) as z:
                    with z.open('word/document.xml') as doc:
                        tree = ET.parse(doc)
                        paragraphs = tree.findall('.//{http://schemas.openxmlformats.org/wordprocessingml/2006/main}t')
                        text = ' '.join(p.text for p in paragraphs if p.text)[:3000]
            except:
                text = ""

    except Exception as e:
        text = ""

    return text.strip()[:2000]  # Cap at 2000 chars to control costs


def describe_document(client, filename, text, file_info):
    """Use Claude to generate a one-line description of the document."""
    if text:
        prompt = f"""You are a legal assistant helping create a privilege log.

Document name: {filename}
File type: {file_info['extension']}
Date modified: {file_info['modified']}

Document content (excerpt):
{text}

Write a single concise sentence (max 20 words) describing what this document is and its legal relevance.
Be specific. Do not use phrases like "This document is" or "This file contains".
Just write the description directly."""
    else:
        prompt = f"""You are a legal assistant helping create a privilege log.

Document name: {filename}
File type: {file_info['extension']}
Date modified: {file_info['modified']}

Based only on the filename and file type, write a single concise sentence (max 15 words) 
describing what this document likely is.
Do not use phrases like "This document is" or "This file contains".
Just write the description directly."""

    message = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=100,
        messages=[{"role": "user", "content": prompt}]
    )

    return message.content[0].text.strip()


def create_excel(entries, output_path, case_name):
    """Create the privilege log Excel spreadsheet."""
    import openpyxl
    from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
    from openpyxl.utils import get_column_letter

    wb = openpyxl.Workbook()
    ws = wb.active
    ws.title = "Privilege Log"

    # Colors
    header_fill = PatternFill(start_color="1a1a2e", end_color="1a1a2e", fill_type="solid")
    alt_fill    = PatternFill(start_color="f5f5f5", end_color="f5f5f5", fill_type="solid")
    border      = Border(
        left=Side(style='thin', color='dddddd'),
        right=Side(style='thin', color='dddddd'),
        top=Side(style='thin', color='dddddd'),
        bottom=Side(style='thin', color='dddddd')
    )

    # Title row
    ws.merge_cells('A1:H1')
    title_cell = ws['A1']
    title = f"PRIVILEGE LOG{' — ' + case_name if case_name else ''}"
    title_cell.value = title
    title_cell.font = Font(name='Arial', bold=True, size=14, color='FFFFFF')
    title_cell.fill = header_fill
    title_cell.alignment = Alignment(horizontal='center', vertical='center')
    ws.row_dimensions[1].height = 30

    # Date row
    ws.merge_cells('A2:H2')
    date_cell = ws['A2']
    date_cell.value = f"Generated: {datetime.now().strftime('%B %d, %Y')}"
    date_cell.font = Font(name='Arial', size=9, color='666666')
    date_cell.alignment = Alignment(horizontal='center')

    # Header row
    headers = ["#", "Filename", "File Type", "Date Modified",
               "File Size", "Description", "Privilege Claimed", "Notes"]
    col_widths = [5, 35, 12, 18, 12, 45, 35, 20]

    for col, (header, width) in enumerate(zip(headers, col_widths), 1):
        cell = ws.cell(row=3, column=col, value=header)
        cell.font = Font(name='Arial', bold=True, size=9, color='FFFFFF')
        cell.fill = header_fill
        cell.alignment = Alignment(horizontal='center', vertical='center', wrap_text=True)
        cell.border = border
        ws.column_dimensions[get_column_letter(col)].width = width

    ws.row_dimensions[3].height = 20

    # Data rows
    for i, entry in enumerate(entries, 1):
        row = i + 3
        fill = alt_fill if i % 2 == 0 else PatternFill()

        values = [
            i,
            entry['filename'],
            entry['file_type'],
            entry['modified'],
            entry['size'],
            entry['description'],
            entry.get('privilege', 'Attorney-Client Privilege'),
            ""
        ]

        for col, value in enumerate(values, 1):
            cell = ws.cell(row=row, column=col, value=value)
            cell.font = Font(name='Arial', size=9)
            cell.alignment = Alignment(vertical='center', wrap_text=True)
            cell.border = border
            if i % 2 == 0:
                cell.fill = alt_fill

        ws.row_dimensions[row].height = 35

    # Freeze panes below header
    ws.freeze_panes = 'A4'

    wb.save(output_path)


def run():
    """Main function."""
    api_key, folder, output_name, case_name = get_input()

    import anthropic
    client = anthropic.Anthropic(api_key=api_key)

    # Get all supported files
    supported = ['.pdf', '.docx', '.txt', '.md', '.csv', '.xlsx', '.png', '.jpg', '.jpeg']
    files = sorted([
        f for f in os.listdir(folder)
        if os.path.splitext(f)[1].lower() in supported
        and not f.startswith('.')
    ])

    if not files:
        print(f"\nNo supported files found in: {folder}")
        print(f"Supported types: {', '.join(supported)}")
        return

    print(f"\nFound {len(files)} files. Generating descriptions with Claude AI...")
    print("(This may take a minute depending on file count)\n")

    entries = []
    ok = bad = 0

    for i, fname in enumerate(files, 1):
        fpath = os.path.join(folder, fname)
        try:
            file_info = get_file_info(fpath)
            text = extract_text(fpath)
            description = describe_document(client, fname, text, file_info)

            entries.append({
                'filename': fname,
                'file_type': file_info['extension'],
                'modified': file_info['modified'],
                'size': file_info['size'],
                'description': description,
                'privilege': 'Attorney-Client Privilege'
            })

            print(f"  ✓ [{i}/{len(files)}] {fname}")
            print(f"       → {description}")
            ok += 1

        except Exception as e:
            print(f"  ✗ [{i}/{len(files)}] {fname} — Error: {e}")
            bad += 1

    # Save Excel file
    output_path = os.path.join(folder, output_name)
    create_excel(entries, output_path, case_name)

    print(f"\nDone! {ok} documents logged, {bad} errors.")
    print(f"Privilege log saved to: {output_path}")
    print("\nTip: Open the Excel file and update the 'Privilege Claimed' and 'Notes' columns as needed.")

    print()
    again = input("Generate another privilege log? (y/n): ").strip().lower()
    if again == "y":
        run()


# ---------------------------------------------------------------
# MAIN
# ---------------------------------------------------------------
if __name__ == '__main__':
    print("=" * 60)
    print("GAVEL — Privilege Log Generator")
    print(f"Running on: {platform.system()} {platform.release()}")
    print("=" * 60)

    if not check_libraries():
        exit(1)

    run()