# app.py - Essay Grader AI with Award List

import gradio as gr
import logging
import os
from datetime import datetime
from docx import Document
from docx.shared import Pt, RGBColor, Inches
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.ns import qn
from docx.oxml import OxmlElement

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

os.makedirs('outputs', exist_ok=True)

try:
    from modules.groq_client import GroqGrader
    from modules.report_generator import PDFReportGenerator
    from modules.firebase_manager import FirebaseManager
    from modules.file_processor import FileProcessor
    from modules.utils import format_results, validate_inputs, get_performance_label
except ImportError as e:
    logger.error(f"Failed to import modules: {e}")
    raise

grader = None
pdf_gen = None
firebase = None

try:
    grader = GroqGrader()
    logger.info("✓ GroqGrader initialized successfully")
except Exception as e:
    logger.error(f"✗ Failed to initialize GroqGrader: {e}")

try:
    pdf_gen = PDFReportGenerator()
    logger.info("✓ PDFReportGenerator initialized successfully")
except Exception as e:
    logger.error(f"✗ Failed to initialize PDFReportGenerator: {e}")

try:
    firebase = FirebaseManager()
    if firebase.enabled:
        logger.info("✓ FirebaseManager initialized successfully")
    else:
        logger.warning("○ Firebase not enabled (optional)")
except Exception as e:
    logger.warning(f"○ Firebase not initialized: {e}")

file_processor = FileProcessor()
current_user = {"id": "demo_user", "subscription": "free"}

# Global storage for class results
class_results = []


def create_word_report(student_name: str, grade_level: int, essay_type: str, result: dict, word_count: int) -> str:
    """Create individual MS Word report"""
    try:
        doc = Document()
        
        title = doc.add_heading('Essay Grading Report', 0)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        date_para = doc.add_paragraph(datetime.now().strftime("%B %d, %Y"))
        date_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        doc.add_paragraph()
        
        doc.add_heading('Student Details', 1)
        doc.add_paragraph(f"Name: {student_name}")
        doc.add_paragraph(f"Grade: {grade_level}")
        doc.add_paragraph(f"Essay Type: {essay_type}")
        doc.add_paragraph(f"Word Count: {word_count}")
        
        doc.add_paragraph()
        
        doc.add_heading('Total Score', 1)
        score = result['total_score']
        label, _, _ = get_performance_label(score)
        score_para = doc.add_paragraph(f"{score}/100 - Grade: {result['grade']} - {label}")
        score_para.runs[0].font.size = Pt(24)
        score_para.runs[0].font.bold = True
        
        doc.add_paragraph()
        
        doc.add_heading('Detailed Scores', 1)
        table = doc.add_table(rows=5, cols=4)
        table.style = 'Light Grid Accent 1'
        
        headers = table.rows[0].cells
        headers[0].text = 'Category'
        headers[1].text = 'Score'
        headers[2].text = 'Maximum'
        headers[3].text = 'Percentage'
        
        categories = [
            ('Content & Arguments', result['scores']['content'], 40),
            ('Organization', result['scores']['organization'], 20),
            ('Language & Style', result['scores']['language'], 20),
            ('Grammar & Spelling', result['scores']['grammar'], 20)
        ]
        
        for i, (name, score_val, max_score) in enumerate(categories, 1):
            row = table.rows[i].cells
            row[0].text = name
            row[1].text = str(score_val)
            row[2].text = str(max_score)
            row[3].text = f"{(score_val/max_score*100):.1f}%"
        
        doc.add_paragraph()
        
        # Structure Analysis
        if 'structure_analysis' in result:
            doc.add_heading('Structure Analysis', 1)
            struct = result['structure_analysis']
            doc.add_paragraph(f"Introduction: {'✓ Present' if struct.get('has_introduction') else '✗ Missing'}")
            doc.add_paragraph(f"Body Paragraphs: {struct.get('body_paragraphs_count', 0)}")
            doc.add_paragraph(f"Transitions: {'✓ Good' if struct.get('has_transitions') else '✗ Weak'}")
            doc.add_paragraph(f"Conclusion: {'✓ Present' if struct.get('has_conclusion') else '✗ Missing'}")
            if struct.get('structure_feedback'):
                doc.add_paragraph(f"Feedback: {struct['structure_feedback']}")
            doc.add_paragraph()
        
        doc.add_heading('Strengths', 1)
        for strength in result.get('strengths', []):
            doc.add_paragraph(strength, style='List Bullet')
        
        doc.add_paragraph()
        
        doc.add_heading('Areas for Improvement', 1)
        for weakness in result.get('weaknesses', []):
            p = doc.add_paragraph(style='List Bullet')
            p.add_run(weakness.get('issue', '')).bold = True
            if 'suggestion' in weakness:
                p.add_run(f"\nSuggestion: {weakness['suggestion']}")
        
        doc.add_paragraph()
        
        doc.add_heading('Overall Feedback', 1)
        doc.add_paragraph(result.get('overall_feedback', ''))
        
        doc.add_paragraph()
        
        footer = doc.add_paragraph("Generated by Essay Grader AI | Developed by: Najaf Ali Sharqi")
        footer.alignment = WD_ALIGN_PARAGRAPH.CENTER
        footer.runs[0].font.size = Pt(9)
        footer.runs[0].font.color.rgb = RGBColor(128, 128, 128)
        
        filename = f"report_{student_name.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
        output_path = os.path.join('outputs', filename)
        doc.save(output_path)
        
        logger.info(f"Word report created: {output_path}")
        return output_path
        
    except Exception as e:
        logger.error(f"Word report creation failed: {e}")
        return None


def create_award_list(institution_name: str, grade_level: int, subject: str) -> str:
    """Create class award list in MS Word"""
    try:
        if not class_results:
            return None
        
        doc = Document()
        
        # Institution Name - Bold, Center, Top
        title = doc.add_heading(institution_name, 0)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        title.runs[0].bold = True
        title.runs[0].font.size = Pt(18)
        
        # Subtitle
        subtitle = doc.add_paragraph(f"Essay Grading Award List")
        subtitle.alignment = WD_ALIGN_PARAGRAPH.CENTER
        subtitle.runs[0].font.size = Pt(14)
        subtitle.runs[0].bold = True
        
        doc.add_paragraph()
        
        # Details
        details = doc.add_paragraph()
        details.add_run(f"Grade: {grade_level}    ").bold = True
        details.add_run(f"Subject: {subject}    ").bold = True
        details.add_run(f"Date: {datetime.now().strftime('%B %d, %Y')}").bold = True
        
        doc.add_paragraph()
        
        # Create table - 5 columns only (removed Feedback Summary)
        num_students = len(class_results)
        table = doc.add_table(rows=num_students + 1, cols=5)
        table.style = 'Light Grid Accent 1'
        
        # Headers
        headers = table.rows[0].cells
        headers[0].text = 'S.No'
        headers[1].text = 'Student Name'
        headers[2].text = 'Score'
        headers[3].text = 'Grade'
        headers[4].text = 'Performance'
        
        # Make headers bold
        for cell in headers:
            cell.paragraphs[0].runs[0].bold = True
            cell.paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # Sort by score (highest first)
        sorted_results = sorted(class_results, key=lambda x: x['result']['total_score'], reverse=True)
        
        # Fill data
        for i, entry in enumerate(sorted_results, 1):
            row = table.rows[i].cells
            result = entry['result']
            score = result['total_score']
            label, _, _ = get_performance_label(score)
            
            row[0].text = str(i)
            row[0].paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            row[1].text = entry['student_name']
            
            row[2].text = f"{score}/100"
            row[2].paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            row[3].text = result['grade']
            row[3].paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
            
            row[4].text = label
            row[4].paragraphs[0].alignment = WD_ALIGN_PARAGRAPH.CENTER
        
        # Add space for signature
        doc.add_paragraph()
        doc.add_paragraph()
        doc.add_paragraph()
        doc.add_paragraph()
        
        signature = doc.add_paragraph()
        signature.add_run("_" * 30)
        signature.add_run("\n")
        signature.add_run("Subject Teacher's Signature").bold = True
        signature.add_run("\n")
        signature.add_run(f"Date: {datetime.now().strftime('%B %d, %Y')}")
        
        # Footer
        doc.add_paragraph()
        footer = doc.add_paragraph("Generated by Essay Grader AI | Developed by: Najaf Ali Sharqi")
        footer.alignment = WD_ALIGN_PARAGRAPH.CENTER
        footer.runs[0].font.size = Pt(8)
        footer.runs[0].font.color.rgb = RGBColor(128, 128, 128)
        
        # Save
        filename = f"award_list_{institution_name.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
        output_path = os.path.join('outputs', filename)
        doc.save(output_path)
        
        logger.info(f"Award list created: {output_path}")
        return output_path
        
    except Exception as e:
        logger.error(f"Award list creation failed: {e}")
        return None


def process_essay(student_name: str, grade_level: int, language: str, essay_type: str, essay_text: str, essay_file, add_to_class: bool) -> tuple:
    try:
        logger.info("=== Starting essay processing ===")
        
        error = validate_inputs(student_name, essay_text, essay_file)
        if error:
            logger.warning(f"Validation failed: {error}")
            return error, None, None, "", ""
        
        if essay_file is not None:
            logger.info(f"Processing file: {essay_file.name}")
            try:
                text = file_processor.extract_text(essay_file.name)
                logger.info(f"Extracted {len(text)} characters from file")
                
                if not text or len(text.strip()) < 10:
                    return "❌ فائل سے متن نہیں نکل سکا۔ براہ کرم:\n\n1. واضح تصویر/scan استعمال کریں\n2. یا متن خود ٹائپ کریں", None, None, "", ""
                
            except Exception as e:
                logger.error(f"File processing error: {e}")
                return f"❌ فائل میں خرابی: {str(e)}", None, None, "", ""
        else:
            text = essay_text
            logger.info(f"Using direct text input: {len(text)} characters")
        
        word_count = len(text.split())
        logger.info(f"Word count: {word_count}")
        
        if word_count < 50:
            return f"❌ مضمون بہت چھوٹا ہے ({word_count} الفاظ)۔ کم از کم 50 الفاظ درکار ہیں۔", None, None, "", ""
        
        if word_count > 5000:
            return f"❌ مضمون بہت لمبا ہے ({word_count} الفاظ)۔ زیادہ سے زیادہ 5000 الفاظ۔", None, None, "", ""
        
        if firebase and firebase.enabled:
            within_limit, used, remaining = firebase.check_monthly_limit(current_user['id'], current_user['subscription'])
            if not within_limit:
                return f"❌ آپ کی ماہانہ حد ({used}) ختم ہو گئی ہے۔", None, None, "", ""
        else:
            used, remaining = 0, 999
        
        if not grader:
            return "❌ Groq API دستیاب نہیں ہے۔ GROQ_API_KEY شامل کریں Settings میں۔", None, None, "", ""
        
        logger.info(f"Calling Groq API for grading...")
        
        lang_code = "urdu" if language == "اردو" else "english"
        
        try:
            result = grader.grade_essay(essay_text=text, language=lang_code, grade_level=grade_level, essay_type=essay_type)
            logger.info(f"✓ Grading successful: {result.get('total_score', 0)}/100")
        except Exception as e:
            logger.error(f"✗ Groq grading error: {e}")
            return f"❌ AI جانچ میں خرابی:\n\n{str(e)}", None, None, "", ""
        
        # Add to class results if checkbox is selected
        if add_to_class:
            class_results.append({
                'student_name': student_name,
                'grade_level': grade_level,
                'essay_type': essay_type,
                'word_count': word_count,
                'result': result
            })
            logger.info(f"Added to class results. Total students: {len(class_results)}")
        
        logger.info("Formatting results...")
        output_md = format_results(student_name=student_name, grade_level=grade_level, essay_type=essay_type, result=result, word_count=word_count)
        
        pdf_path = None
        word_path = None
        
        if pdf_gen:
            try:
                logger.info("Creating PDF report...")
                pdf_filename = f"report_{student_name.replace(' ', '_')}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
                pdf_path = os.path.join('outputs', pdf_filename)
                pdf_gen.create_report(student_name=student_name, grade_level=grade_level, essay_type=essay_type, grading_result=result, output_path=pdf_path)
                logger.info(f"✓ PDF created: {pdf_path}")
            except Exception as e:
                logger.error(f"✗ PDF generation error: {e}")
        
        try:
            logger.info("Creating Word report...")
            word_path = create_word_report(student_name, grade_level, essay_type, result, word_count)
            if word_path:
                logger.info(f"✓ Word created: {word_path}")
        except Exception as e:
            logger.error(f"✗ Word generation error: {e}")
        
        if firebase and firebase.enabled:
            try:
                essay_data = {
                    'student_name': student_name,
                    'grade_level': grade_level,
                    'essay_type': essay_type,
                    'language': lang_code,
                    'word_count': word_count,
                    'text': text[:1000]
                }
                firebase.save_essay_result(user_id=current_user['id'], essay_data=essay_data, grading_result=result)
                logger.info("✓ Essay saved to Firebase")
            except Exception as e:
                logger.error(f"✗ Firebase save error: {e}")
        
        limit_msg = f"📊 استعمال: {used + 1}, باقی: {remaining - 1}"
        class_msg = f"📚 کلاس میں طلبا: {len(class_results)}" if add_to_class else ""
        
        logger.info("=== Essay processing completed successfully ===")
        return output_md, pdf_path, word_path, limit_msg, class_msg
        
    except Exception as e:
        logger.exception("✗ Unexpected error in process_essay")
        return f"❌ غیر متوقع خرابی:\n\n{str(e)}", None, None, "", ""


def generate_award_list_handler(institution_name: str, grade_level: int, subject: str):
    """Handle award list generation"""
    try:
        if not institution_name or not institution_name.strip():
            return "❌ ادارے کا نام درکار ہے | Institution name required", None
        
        if not subject or not subject.strip():
            return "❌ مضمون کا نام درکار ہے | Subject name required", None
        
        if not class_results:
            return "❌ کوئی طالب علم شامل نہیں۔ پہلے مضامین جانچیں اور 'Add to Class Award List' چیک کریں۔", None
        
        award_path = create_award_list(institution_name, grade_level, subject)
        
        if award_path:
            return f"✅ Award List بن گئی! کل طلبا: {len(class_results)}", award_path
        else:
            return "❌ Award List بنانے میں خرابی", None
        
    except Exception as e:
        logger.error(f"Award list generation error: {e}")
        return f"❌ خرابی: {str(e)}", None


def clear_form():
    """Clear all form inputs"""
    return "", None, "", ""


def reset_class_list():
    """Reset class results"""
    global class_results
    class_results = []
    return f"✅ کلاس کی فہرست صاف ہو گئی | Class list cleared"


app = gr.Blocks(title="Essay Grader AI")

with app:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px; margin-bottom: 20px;">
        <h1 style="margin: 5px 0;">📝 Essay Grader AI</h1>
        <p style="margin: 5px 0; font-size: 14px;">Professional AI-Powered Essay Grading System with OCR & Award List</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📋 مضمون کی تفصیلات | Essay Details")
            
            student_name = gr.Textbox(label="طالب علم کا نام | Student Name", placeholder="مثال: محمد احمد")
            
            with gr.Row():
                grade_level = gr.Dropdown(choices=list(range(5, 13)), label="جماعت | Grade", value=8)
                language = gr.Radio(choices=["اردو", "English"], label="زبان | Language", value="اردو")
            
            essay_type = gr.Dropdown(
                choices=["تشریحی (Descriptive)", "استدلالی (Argumentative)", "تجزیاتی (Analytical)", "تخلیقی (Creative)"],
                label="مضمون کی قسم | Essay Type",
                value="تشریحی (Descriptive)"
            )
            
            add_to_class_checkbox = gr.Checkbox(
                label="✅ Add to Class Award List | کلاس کی فہرست میں شامل کریں",
                value=False,
                info="Check this to include in class award list"
            )
            
            gr.Markdown("### 📄 مضمون | Essay")
            
            essay_file = gr.File(
                label="فائل اپ لوڈ | Upload File (.txt, .docx, .pdf - handwritten OK!)", 
                file_types=['.txt', '.docx', '.pdf'], 
                type="filepath"
            )
            
            gr.Markdown("""
**یا براہ راست لکھیں | Or write directly:**

💡 **نوٹ**: Scanned/handwritten PDFs: 30-60s processing time
            """)
            
            essay_text = gr.Textbox(
                label="مضمون یہاں لکھیں | Write Essay Here", 
                lines=8, 
                placeholder="اپنا مضمون یہاں لکھیں یا پیسٹ کریں..."
            )
            
            with gr.Row():
                submit_btn = gr.Button("🔍 جانچیں | Grade Essay", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ صاف کریں | Clear", variant="secondary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("""
### خوش آمدید! Welcome!

**خصوصیات | Features:**
- ⚡ 20-30 سیکنڈ میں مکمل تجزیہ
- 📊 تفصیلی سکور کارڈ (Consistent)
- 🏗️ ساخت کا مکمل تجزیہ (Structure Analysis)
- 💡 عملی بہتری کے مشورے
- 📄 پیشہ ورانہ PDF رپورٹ
- 📝 MS Word فائل ڈاؤن لوڈ
- 🌐 اردو اور انگریزی
- 🖼️ Handwritten/Scanned PDF Support (OCR)
- 🏆 Class Award List Generator

**Performance Scale:**
- 90+: Exceptional 🏆
- 80-89: Excellent ⭐
- 70-79: Very Good 👍
- 60-69: Good ✓
- 50-59: Pass 📝
- <50: Needs Improvement 💪

---
            """)
            
            gr.Markdown("### 📊 نتائج | Results")
            
            limit_display = gr.Markdown("")
            class_status = gr.Markdown("")
            
            results_output = gr.Markdown(value="نتائج یہاں ظاہر ہوں گے...\n\nResults will appear here...")
            
            with gr.Row():
                pdf_output = gr.File(label="📥 PDF رپورٹ | PDF Report")
                word_output = gr.File(label="📥 Word رپورٹ | Word Report")
    
    # Award List Section
    gr.Markdown("---")
    gr.Markdown("## 🏆 Class Award List Generator")
    
    with gr.Row():
        with gr.Column():
            institution_name = gr.Textbox(
                label="ادارے کا نام | Institution Name",
                placeholder="مثال: Government High School, Bannu",
                value=""
            )
            
            with gr.Row():
                award_grade = gr.Dropdown(
                    choices=list(range(5, 13)),
                    label="جماعت | Grade",
                    value=8
                )
                subject_name = gr.Textbox(
                    label="مضمون | Subject",
                    placeholder="مثال: English Composition / اردو تحریر",
                    value="Essay Writing"
                )
            
            with gr.Row():
                generate_award_btn = gr.Button("🏆 Generate Award List", variant="primary", size="lg")
                reset_class_btn = gr.Button("🔄 Reset Class List", variant="secondary")
        
        with gr.Column():
            award_status = gr.Markdown("")
            award_output = gr.File(label="📥 Download Award List")
    
    # Event handlers
    submit_btn.click(
        fn=process_essay,
        inputs=[student_name, grade_level, language, essay_type, essay_text, essay_file, add_to_class_checkbox],
        outputs=[results_output, pdf_output, word_output, limit_display, class_status]
    )
    
    clear_btn.click(
        fn=clear_form,
        inputs=[],
        outputs=[student_name, essay_file, essay_text, results_output]
    )
    
    generate_award_btn.click(
        fn=generate_award_list_handler,
        inputs=[institution_name, award_grade, subject_name],
        outputs=[award_status, award_output]
    )
    
    reset_class_btn.click(
        fn=reset_class_list,
        inputs=[],
        outputs=[award_status]
    )
    
    gr.Markdown("""
    ---
    <div style="text-align: center; padding: 15px; background-color: #f5f5f5; border-radius: 10px;">
        <p style="margin: 5px 0;"><strong>AI Model:</strong> Groq (Llama-3.1-70B) | <strong>Framework:</strong> Gradio | <strong>OCR:</strong> Tesseract</p>
        <p style="margin: 5px 0; color: #666;">Powered by Groq AI • Made with ❤️ for Teachers and Students</p>
        <p style="margin: 5px 0; color: #1E40AF; font-weight: bold;">Developed by: Najaf Ali Sharqi</p>
        <p style="margin: 5px 0; color: #666; font-size: 11px;">© 2024 Essay Grader AI • GB Gen AI Project</p>
    </div>
    """)

if __name__ == "__main__":
    logger.info("Starting Essay Grader AI application")
    app.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)