#!/usr/bin/env python3
"""
Tommy's Homework Annotator v2
Converts submitted writing files to annotated PNG images with teacher-style feedback.
New approach: Clean text rendering + gray summary box at bottom with rubric breakdown.
"""

import sys
import json
import os
import subprocess
import tempfile
import csv
from pathlib import Path
from datetime import datetime
from PIL import Image, ImageDraw, ImageFont


def convert_to_png(input_file):
    """
    Convert various file formats to PNG image(s).
    Returns list of PNG file paths.
    """
    input_path = Path(input_file).resolve()
    temp_dir = tempfile.gettempdir()
    
    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_file}")
    
    suffix = input_path.suffix.lower()
    png_files = []
    
    try:
        if suffix == ".txt":
            # TXT → Direct conversion via text rendering
            png_files = [render_text_to_png(input_path, temp_dir)]
        
        elif suffix == ".pages":
            # Pages → PDF → PNG
            pdf_path = f"{temp_dir}/hw_converted.pdf"
            subprocess.run(
                ["textutil", "-convert", "pdf", str(input_path), "-output", pdf_path],
                check=True,
                capture_output=True
            )
            png_files = pdf_to_png(pdf_path, temp_dir)
        
        elif suffix == ".pdf":
            # PDF → PNG
            png_files = pdf_to_png(str(input_path), temp_dir)
        
        elif suffix == ".docx":
            # DOCX → Extract text → PNG
            try:
                import docx
                doc = docx.Document(str(input_path))
                text = "\n".join([p.text for p in doc.paragraphs])
                png_files = [render_text_to_png_from_text(text, temp_dir)]
            except ImportError:
                # Fallback: use textutil if available
                pdf_path = f"{temp_dir}/hw_converted.pdf"
                try:
                    subprocess.run(
                        ["textutil", "-convert", "pdf", str(input_path), "-output", pdf_path],
                        check=True,
                        capture_output=True
                    )
                    png_files = pdf_to_png(pdf_path, temp_dir)
                except:
                    # Last resort: render a placeholder
                    png_files = [render_text_to_png_from_text(f"[DOCX file: {input_path.name}]", temp_dir)]
        
        else:
            raise ValueError(f"Unsupported file format: {suffix}")
        
        return png_files
    
    except Exception as e:
        print(f"Error converting file: {e}", file=sys.stderr)
        raise


def render_text_to_png(txt_path, output_dir):
    """Render a .txt file to PNG with clean typography."""
    with open(txt_path, "r") as f:
        text = f.read()
    return render_text_to_png_from_text(text, output_dir)


def render_text_to_png_from_text(text, output_dir, reserve_bottom=200):
    """Render text content to PNG image, reserving space at bottom for summary box."""
    # Page dimensions
    page_width = 850
    page_height = 1100
    margin = 50
    line_height = 24
    
    # Create image with white background
    img = Image.new("RGB", (page_width, page_height), color="white")
    draw = ImageDraw.Draw(img)
    
    # Try to use a system font; fall back to default
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 14)
    except:
        font = ImageFont.load_default()
    
    # Draw text with line wrapping, reserving space at bottom
    y = margin
    x = margin
    max_width = page_width - 2 * margin
    max_y = page_height - reserve_bottom  # Stop before summary box area
    
    # Simple word wrap
    words = text.split()
    line = ""
    for word in words:
        test_line = f"{line} {word}".strip()
        bbox = draw.textbbox((0, 0), test_line, font=font)
        line_width = bbox[2] - bbox[0]
        
        if line_width > max_width and line:
            draw.text((x, y), line, fill="black", font=font)
            y += line_height
            line = word
            if y > max_y:
                break
        else:
            line = test_line
    
    if line and y <= max_y:
        draw.text((x, y), line, fill="black", font=font)
    
    # Save to temp location
    png_path = os.path.join(output_dir, "hw_rendered.png")
    img.save(png_path)
    return png_path


def pdf_to_png(pdf_path, output_dir):
    """Convert PDF to PNG image(s) using sips or ghostscript."""
    png_files = []
    
    # Try sips first (native macOS)
    try:
        png_path = os.path.join(output_dir, "hw_page.png")
        result = subprocess.run(
            ["sips", "-s", "format", "png", pdf_path, "--out", png_path],
            capture_output=True,
            text=True
        )
        if result.returncode == 0 and os.path.exists(png_path):
            png_files.append(png_path)
            return png_files
    except:
        pass
    
    # Fallback to ghostscript
    try:
        png_pattern = os.path.join(output_dir, "hw_page_%d.png")
        subprocess.run(
            [
                "gs", "-dNOPAUSE", "-dBATCH", "-sDEVICE=png256",
                "-r150", f"-sOutputFile={png_pattern}", pdf_path
            ],
            check=True,
            capture_output=True
        )
        # Find generated PNG files
        import glob
        pattern = os.path.join(output_dir, "hw_page_*.png")
        png_files = sorted(glob.glob(pattern))
        if png_files:
            return png_files
    except:
        pass
    
    # If both fail, render a placeholder
    placeholder = render_text_to_png_from_text(f"[PDF: {pdf_path}]", output_dir)
    return [placeholder]


def add_summary_box(image_path, grade_data):
    """
    Add a gray summary box at the BOTTOM of the image with rubric breakdown.
    Returns annotated PIL Image object.
    """
    img = Image.open(image_path).convert("RGB")
    draw = ImageDraw.Draw(img)
    
    # Get fonts
    try:
        font_regular = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 12)
        font_bold = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 14)
        font_grade = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 48)
    except:
        font_regular = ImageFont.load_default()
        font_bold = font_regular
        font_grade = font_regular
    
    width, height = img.size
    
    # Summary box dimensions
    box_margin = 30
    box_y_start = height - 200  # Start 200px from bottom
    box_height = 170
    box_x1 = box_margin
    box_x2 = width - box_margin
    box_y1 = box_y_start
    box_y2 = box_y_start + box_height
    
    # Draw gray background box
    draw.rectangle(
        [box_x1, box_y1, box_x2, box_y2],
        fill="#E8E8E8",
        outline="#AAAAAA",
        width=2
    )
    
    # Grade (large red text, centered at top of box)
    grade = grade_data.get("score", "?")
    
    # Measure grade text to center it
    grade_bbox = draw.textbbox((0, 0), grade, font=font_grade)
    grade_width = grade_bbox[2] - grade_bbox[0]
    grade_x = box_x1 + (box_x2 - box_x1 - grade_width) // 2
    grade_y = box_y1 + 10
    
    draw.text((grade_x, grade_y), grade, fill="#FF0000", font=font_grade)
    
    # Rubric breakdown (beneath grade)
    rubric_y = grade_y + 55
    rubric = grade_data.get("rubric_breakdown", "")
    
    if rubric:
        # Format: "Thesis 22/25 | Organization 24/25 | Evidence 25/25"
        rubric_lines = rubric.split(" | ")
        rubric_text = " | ".join(rubric_lines[:3])  # Limit to 3 items per line
        draw.text((box_x1 + 10, rubric_y), rubric_text, fill="#333333", font=font_regular)
    
    # Overall feedback (1-2 lines)
    feedback_y = rubric_y + 25
    feedback = grade_data.get("overall_feedback", "")
    
    if feedback:
        # Wrap feedback to fit
        feedback_lines = []
        words = feedback.split()
        current_line = ""
        max_feedback_width = box_x2 - box_x1 - 20
        
        for word in words:
            test_line = f"{current_line} {word}".strip()
            bbox = draw.textbbox((0, 0), test_line, font=font_regular)
            line_width = bbox[2] - bbox[0]
            
            if line_width > max_feedback_width and current_line:
                feedback_lines.append(current_line)
                current_line = word
            else:
                current_line = test_line
        
        if current_line:
            feedback_lines.append(current_line)
        
        for idx, line in enumerate(feedback_lines[:2]):  # Max 2 lines
            draw.text((box_x1 + 10, feedback_y + idx * 18), line, fill="#333333", font=font_regular)
    
    # Improve section (numbered list, max 3 points)
    improve_y = feedback_y + 50
    improvements = grade_data.get("improvements", [])
    
    if improvements:
        draw.text((box_x1 + 10, improve_y), "Improve:", fill="#333333", font=font_bold)
        
        for idx, improvement in enumerate(improvements[:3]):
            improvement_text = f"{idx + 1}. {improvement[:50]}"  # Truncate long text
            draw.text((box_x1 + 25, improve_y + 18 + idx * 16), improvement_text, fill="#333333", font=font_regular)
    
    return img


def log_grade_to_csv(grade_data, assignment_name):
    """
    Log grade to CSV file in ~/.openclaw/workspace/academic-records/[Subject]-grades.csv
    Creates the directory if it doesn't exist.
    """
    # Determine subject from grade_data
    subject = grade_data.get("subject", "General")
    
    # Create directory if needed (in Tommy's workspace, not Documents)
    records_dir = Path.home() / ".openclaw" / "workspace" / "academic-records"
    records_dir.mkdir(parents=True, exist_ok=True)
    
    # CSV file path
    csv_path = records_dir / f"{subject}-grades.csv"
    
    # Prepare row data
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    letter_grade = grade_data.get("score", "?")
    percentage = grade_data.get("percentage", 0)
    
    # Get individual rubric scores (if available)
    # For writing: thesis, organization, evidence, grammar, effort
    # For computation: correctness, process, effort
    scores = grade_data.get("scores", {})
    
    if subject.lower() in ["math", "science", "computation"]:
        # Computation rubric
        row = {
            "date": timestamp,
            "assignment_name": assignment_name,
            "subject": subject,
            "letter_grade": letter_grade,
            "percentage": percentage,
            "correctness_score": scores.get("correctness", ""),
            "process_score": scores.get("process", ""),
            "effort_score": scores.get("effort", ""),
            "tommy_note": grade_data.get("overall_feedback", "")[:100]
        }
    else:
        # Writing/Analysis rubric
        row = {
            "date": timestamp,
            "assignment_name": assignment_name,
            "subject": subject,
            "letter_grade": letter_grade,
            "percentage": percentage,
            "thesis_score": scores.get("thesis", ""),
            "organization_score": scores.get("organization", ""),
            "evidence_score": scores.get("evidence", ""),
            "grammar_score": scores.get("grammar", ""),
            "effort_score": scores.get("effort", ""),
            "tommy_note": grade_data.get("overall_feedback", "")[:100]
        }
    
    # Write to CSV (append if exists, create if not)
    file_exists = csv_path.exists()
    
    with open(csv_path, "a", newline="") as f:
        fieldnames = list(row.keys())
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        
        if not file_exists:
            writer.writeheader()
        
        writer.writerow(row)
    
    return str(csv_path)


def save_grade_report(output_dir, filename, grade_data):
    """Save a markdown grade report."""
    timestamp = datetime.now().strftime("%Y-%m-%d")
    report_name = f"{timestamp}-{Path(filename).stem}-grade.md"
    report_path = os.path.join(output_dir, report_name)
    
    score = grade_data.get("score", "?")
    percentage = grade_data.get("percentage", 0)
    overall = grade_data.get("overall_feedback", "")
    improvements = grade_data.get("improvements", [])
    subject = grade_data.get("subject", "General")
    scores = grade_data.get("scores", {})
    
    # Build rubric breakdown
    if subject.lower() in ["math", "science", "computation"]:
        rubric_section = f"""## Rubric Breakdown
| Criterion | Score |
|-----------|-------|
| Correctness | {scores.get("correctness", "?")} |
| Process/Showing Work | {scores.get("process", "?")} |
| Effort/Completeness | {scores.get("effort", "?")} |
"""
    else:
        rubric_section = f"""## Rubric Breakdown
| Criterion | Score |
|-----------|-------|
| Thesis/Main Point | {scores.get("thesis", "?")} |
| Organization | {scores.get("organization", "?")} |
| Evidence/Reasoning | {scores.get("evidence", "?")} |
| Grammar/Mechanics | {scores.get("grammar", "?")} |
| Effort/Depth | {scores.get("effort", "?")} |
"""
    
    report = f"""# Grade Report
- Date: {timestamp}
- File: {filename}
- Subject: {subject}
- Grade: {score} ({percentage}%)

## Summary
{overall}

{rubric_section}

## Areas for Improvement
"""
    
    for idx, improvement in enumerate(improvements, 1):
        report += f"{idx}. {improvement}\n"
    
    with open(report_path, "w") as f:
        f.write(report)
    
    return report_path


def main():
    """Main entry point."""
    if len(sys.argv) < 3:
        print("Usage: annotate.py <input_file> <grade_json>")
        print("Example: annotate.py essay.txt '{\"score\":\"B+\",\"percentage\":88,...}'")
        sys.exit(1)
    
    input_file = sys.argv[1]
    grade_json = sys.argv[2]
    
    # Parse grade data
    try:
        grade_data = json.loads(grade_json)
    except json.JSONDecodeError:
        print(f"Error: Invalid JSON in grade data", file=sys.stderr)
        sys.exit(1)
    
    # Create output directory if needed
    output_dir = os.path.expanduser("~/Desktop/Graded-Work")
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # Convert file to PNG(s)
        print(f"Converting {input_file} to PNG...")
        png_files = convert_to_png(input_file)
        
        if not png_files:
            print("Error: No PNG files generated", file=sys.stderr)
            sys.exit(1)
        
        # Add summary box to first page
        png_path = png_files[0]
        print(f"Adding summary box...")
        annotated_img = add_summary_box(png_path, grade_data)
        
        # Save annotated PNG
        timestamp = datetime.now().strftime("%Y-%m-%d")
        filename = Path(input_file).stem
        output_name = f"{timestamp}-{filename}-annotated.png"
        output_path = os.path.join(output_dir, output_name)
        
        annotated_img.save(output_path)
        print(f"✓ Saved annotated image: {output_path}")
        
        # Save grade report
        report_path = save_grade_report(output_dir, Path(input_file).name, grade_data)
        print(f"✓ Saved grade report: {report_path}")
        
        # Log grade to CSV
        csv_path = log_grade_to_csv(grade_data, Path(input_file).stem)
        print(f"✓ Logged grade to CSV: {csv_path}")
        
        # Return paths for shell integration
        print(f"ANNOTATED_PNG={output_path}")
        print(f"GRADE_REPORT={report_path}")
        print(f"GRADE_CSV={csv_path}")
        
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()