from flask import Flask, render_template, request, jsonify, send_file
import google.generativeai as genai
import base64
import logging
from weasyprint import HTML
import os
from datetime import datetime
import tempfile
from io import BytesIO
import jinja2
from dotenv import load_dotenv
from tenacity import retry, stop_after_attempt, wait_exponential

app = Flask(__name__)

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

# Configure Gemini API with error handling
api_key = os.getenv('GOOGLE_API_KEY')
if not api_key:
    error_msg = ("No Google API key found. "
                 "For Hugging Face deployment, please add GOOGLE_API_KEY "
                 "in your Space's Settings -> Repository Secrets")
    logger.error(error_msg)
    raise ValueError(error_msg)

try:
    genai.configure(api_key=api_key)
    
    # Configure the model with safety settings
    generation_config = {
        "temperature": 0.9,
        "top_p": 1,
        "top_k": 1,
        "max_output_tokens": 2048,
    }

    safety_settings = [
        {
            "category": "HARM_CATEGORY_HARASSMENT",
            "threshold": "BLOCK_MEDIUM_AND_ABOVE"
        },
        {
            "category": "HARM_CATEGORY_HATE_SPEECH",
            "threshold": "BLOCK_MEDIUM_AND_ABOVE"
        },
        {
            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
            "threshold": "BLOCK_MEDIUM_AND_ABOVE"
        },
        {
            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
            "threshold": "BLOCK_MEDIUM_AND_ABOVE"
        }
    ]

    model = genai.GenerativeModel(
        model_name="gemini-1.5-flash",
        generation_config=generation_config,
        safety_settings=safety_settings
    )
    
    # Test the connection
    response = model.generate_content("Test connection")
    logger.info("Successfully configured Gemini API")
except Exception as e:
    logger.error(f"Failed to configure Gemini API: {str(e)}")
    raise

# Updated prompt for dual-format report
prompt = """You are a professional campus facility inspector with over 15 years of experience in infrastructure assessment in India. Analyze the campus with total area of ${college_area} acres. Generate two reports based on the provided campus images:
+ The campus includes ${ground_count} grounds with a total area of ${ground_area} acres.
  
  REPORT 1: EXECUTIVE SUMMARY TABLES
  
  Table 1: Campus Overview
  | Aspect | Grade | Key Observations | Priority Level |
  |--------|-------|-----------------|----------------|
  | Overall Infrastructure | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Buildings | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Roads & Parking | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Sports Grounds (${ground_area} acres) | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Canteens | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Entry/Exit & Security | (A+/A/B+/B/C) | • Key points | High/Medium/Low |
  | Environmental | (A+/A/B+/B/C) | • Key points | High/Medium/Low |

Table 2: Critical Issues and Solutions
| Area | Issue | Proposed Solution/Measures |
|------|-------|---------------------------|
| Area 1 | Description of issue | • Detailed solution steps |
| Area 2 | Description of issue | • Detailed solution steps |

REPORT 2: DETAILED ASSESSMENT

1. Overall Campus Analysis
   A. Infrastructure Overview
      - General campus layout and planning
      - Common areas and circulation
      - Campus-wide systems (drainage, lighting)
      - Shared facilities condition
   
   B. Safety & Security Assessment
      - Boundary security
      - Emergency systems
      - Lighting and surveillance
      - Fire safety measures
   
   C. Environmental Analysis
      - Green spaces and landscaping
      - Water management
      - Waste management
      - Natural lighting and ventilation
   
   D. Accessibility & Connectivity
      - Internal roads and pathways
      - Parking facilities
      - Emergency access
      - Campus connectivity

2. Area-wise Assessment
   A. Buildings
      - Overall condition
      - Key features
      - Notable issues
      - Maintenance status

   B. Roads & Parking
      - Surface condition
      - Traffic flow
      - Parking adequacy
      - Safety features

   C. Sports Facilities
      - Ground conditions
      - Equipment status
      - Safety measures
      - Maintenance level

   D. Canteens
      - Structure condition
      - Hygiene standards
      - Ventilation
      - Seating capacity

   E. Entry/Exit Points
      - Security measures
      - Traffic management
      - Access control
      - Emergency preparedness

3. Campus Strengths
   A. Infrastructure Strengths
      - Notable features
      - Well-maintained areas
      - Effective systems
      - Best practices observed

   B. Enhancement Potential
      - Areas showing excellence
      - Opportunities for showcase
      - Positive aspects to build upon
      - Innovative features

4. Areas of Concern & Recommendations
   A. Critical Issues
      - Infrastructure gaps
      - Safety concerns
      - Maintenance needs
      - Operational challenges

   B. Improvement Measures
      - Specific solutions
      - Practical steps
      - Enhancement strategies
      - Preventive measures

5. Final Assessment
   A. Overall Grade: [A+/A/B+/B/C]
      Brief justification of the grade based on:
      - Infrastructure quality
      - Maintenance standards
      - Safety measures
      - Environmental aspects

   B. Concluding Remarks
      - Key takeaways
      - Critical focus areas
      - Positive highlights
      - Path forward

Please focus on significant issues only and ignore minor cosmetic concerns. Consider local weather patterns (monsoon, summer) and regional building practices. Present information in clear, concise formats."""

@app.route('/')
def index():
    return render_template('index.html')

def calculate_grade(grades, weights):
    total_weight = sum(weights.values())
    weighted_score = sum(grades[aspect] * weights[aspect] for aspect in grades)
    average_score = weighted_score / total_weight
    # Convert average score to a grade
    if average_score >= 90:
        return 'A+'
    elif average_score >= 80:
        return 'A'
    elif average_score >= 70:
        return 'B+'
    elif average_score >= 60:
        return 'B'
    else:
        return 'C'

def extract_grades(report_text):
    # Dummy implementation for extracting grades from report text
    # This should be replaced with actual logic to parse the report
    return {
        'Overall Infrastructure': 85,
        'Buildings': 80,
        'Roads & Parking': 75,
        'Sports Facilities': 70,
        'Canteens': 65,
        'Entry/Exit & Security': 80,
        'Environmental': 90
    }

def extract_priority_distribution(report_text):
    """Extract priority distribution from the report text"""
    try:
        # Count priority levels from the Overview table
        priorities = {
            'High': 0,
            'Medium': 0,
            'Low': 0
        }
        
        # Look for priority levels in the Overview table
        lines = report_text.split('\n')
        for line in lines:
            if '|' in line:  # Table row
                if 'High' in line:
                    priorities['High'] += 1
                elif 'Medium' in line:
                    priorities['Medium'] += 1
                elif 'Low' in line:
                    priorities['Low'] += 1
                    
        return [priorities['High'], priorities['Medium'], priorities['Low']]
    except Exception as e:
        logger.error(f"Error extracting priority distribution: {str(e)}")
        return [30, 45, 25]  # Default values if extraction fails

def extract_area_performance(report_text):
    """Extract performance scores for different areas"""
    try:
        # Extract scores from the Overview table
        areas = {
            'Buildings': 0,
            'Roads & Parking': 0,
            'Sports Facilities': 0,
            'Canteens': 0,
            'Security': 0,
            'Environmental': 0
        }
        
        lines = report_text.split('\n')
        for line in lines:
            if '|' in line:  # Table row
                parts = line.split('|')
                if len(parts) >= 2:
                    area = parts[1].strip()
                    if area in areas:
                        # Convert grade to numeric score
                        grade = parts[2].strip() if len(parts) > 2 else ''
                        score = {
                            'A+': 95,
                            'A': 85,
                            'B+': 75,
                            'B': 65,
                            'C': 55
                        }.get(grade, 70)
                        areas[area] = score
                        
        return list(areas.values())
    except Exception as e:
        logger.error(f"Error extracting area performance: {str(e)}")
        return [85, 75, 70, 80, 90, 85]  # Default values if extraction fails

def extract_maintenance_status(report_text):
    """Extract maintenance status distribution"""
    try:
        status = {
            'Well Maintained': 0,
            'Needs Attention': 0,
            'Critical': 0
        }
        
        # Count maintenance status mentions in the report
        well_maintained = report_text.lower().count('well maintained')
        needs_attention = report_text.lower().count('needs attention') + report_text.lower().count('requires attention')
        critical = report_text.lower().count('critical') + report_text.lower().count('urgent')
        
        total = well_maintained + needs_attention + critical or 1  # Avoid division by zero
        
        return [
            int(well_maintained * 100 / total),
            int(needs_attention * 100 / total),
            int(critical * 100 / total)
        ]
    except Exception as e:
        logger.error(f"Error extracting maintenance status: {str(e)}")
        return [60, 30, 10]  # Default values if extraction fails

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def generate_content_with_retry(prompt, images):
    """Generate content with retry logic"""
    try:
        response = model.generate_content(
            [prompt] + images,
            generation_config=genai.types.GenerationConfig(
                # Remove timeout parameter
                # Other generation config parameters can be added here if needed
            )
        )
        return response
    except Exception as e:
        logger.error(f"Error generating content: {str(e)}")
        raise

@app.route('/generate_report', methods=['POST'])
def generate_report():
    try:
        if not api_key:
            raise ValueError("Google API key not configured")

        data = request.json
        images = data.get('images', [])
        basic_info = data.get('basicInfo', {})
        
        # Add file size validation
        MAX_FILE_SIZE = 4 * 1024 * 1024  # 4MB
        for image_data in images:
            if len(base64.b64decode(image_data['data'].split(',')[1])) > MAX_FILE_SIZE:
                return jsonify({'error': 'Image file size exceeds 4MB limit'}), 400
        
        # Create image context with numbering
        image_contexts = []
        all_images = []
        image_count = 1
        
        for image_data in images:
            if image_data['data'].startswith('data:image'):
                image_data['data'] = image_data['data'].split(',')[1]
            image_bytes = base64.b64decode(image_data['data'])
            
            # Store image data for the report
            image_context = {
                'number': image_count,
                'category': image_data['category'],
                'data': image_bytes,
                'mime_type': 'image/jpeg'
            }
            image_contexts.append(image_context)
            
            # Add to list for API
            all_images.append({
                "mime_type": "image/jpeg",
                "data": image_bytes
            })
            
            image_count += 1

        # Update prompt with actual values
        contextualized_prompt = prompt.replace('${college_area}', str(basic_info.get('collegeArea', '')))
        contextualized_prompt = contextualized_prompt.replace('${building_count}', str(basic_info.get('buildingCount', '')))
        contextualized_prompt = contextualized_prompt.replace('${parking_area}', str(basic_info.get('parkingCount', '')))
        contextualized_prompt = contextualized_prompt.replace('${canteen_count}', str(basic_info.get('canteenCount', '')))
        contextualized_prompt = contextualized_prompt.replace('${ground_count}', str(basic_info.get('groundCount', '')))
        contextualized_prompt = contextualized_prompt.replace('${ground_area}', str(basic_info.get('groundArea', '')))
        contextualized_prompt = contextualized_prompt.replace('${gate_count}', str(basic_info.get('gateCount', '')))
        
        # Generate report with image references
        response = generate_content_with_retry(
            contextualized_prompt,
            all_images
        )
        report_text = response.text
        
        # Extract all metrics
        grades = extract_grades(report_text)
        priority_distribution = extract_priority_distribution(report_text)
        area_performance = extract_area_performance(report_text)
        maintenance_status = extract_maintenance_status(report_text)
        
        # Calculate overall grade
        weights = {
            'Overall Infrastructure': 0.2,
            'Buildings': 0.2,
            'Roads & Parking': 0.15,
            'Sports Facilities': 0.1,
            'Canteens': 0.1,
            'Entry/Exit & Security': 0.15,
            'Environmental': 0.1
        }
        overall_grade = calculate_grade(grades, weights)

        # Add metrics to the response
        return jsonify({
            'report': report_text,
            'overallGrade': overall_grade,
            'metrics': {
                'priorityDistribution': priority_distribution,
                'areaPerformance': area_performance,
                'maintenanceStatus': maintenance_status
            },
            'images': [{
                'number': img['number'],
                'category': img['category'],
                'data': base64.b64encode(img['data']).decode('utf-8')
            } for img in image_contexts]
        })
        
    except ValueError as ve:
        logger.error(f"Validation error: {str(ve)}")
        return jsonify({'error': str(ve)}), 400
    except Exception as e:
        logger.error(f"Error generating report: {str(e)}")
        return jsonify({'error': 'Internal server error occurred'}), 500

@app.route('/download_pdf', methods=['POST'])
def download_pdf():
    pdf_buffer = None
    try:
        logger.debug("Received PDF download request")
        html_content = request.json.get('html')
        
        if not html_content:
            raise ValueError("No HTML content provided")

        logger.debug("Converting HTML to PDF")
        # Create PDF in memory
        pdf_buffer = BytesIO()
        HTML(string=html_content).write_pdf(pdf_buffer)
        pdf_buffer.seek(0)
        
        logger.debug("Sending PDF file")
        # Create a copy of the buffer contents
        pdf_data = pdf_buffer.getvalue()
        
        # Close the original buffer
        if pdf_buffer:
            pdf_buffer.close()
            
        # Create a new buffer with the copied data
        return_buffer = BytesIO(pdf_data)
        
        return send_file(
            return_buffer,
            mimetype='application/pdf',
            as_attachment=True,
            download_name=f'campus-inspection-report-{datetime.now().strftime("%Y%m%d")}.pdf'
        )

    except Exception as e:
        logger.error(f"Error generating PDF: {str(e)}")
        return jsonify({'error': str(e)}), 500
    finally:
        # Clean up the original buffer if it exists
        if pdf_buffer:
            pdf_buffer.close()

def generate_report(findings, inspector_name, location, weather):
    # Format timestamp
    timestamp = datetime.now().strftime("%B %d, %Y at %I:%M %p")
    
    # Ensure findings have all required fields and proper formatting
    formatted_findings = []
    for finding in findings:
        formatted_finding = {
            'title': finding.get('title', 'Untitled Finding'),
            'description': finding.get('description', 'No description provided'),
            'severity': finding.get('severity', 'Medium').capitalize(),
            'recommendation': finding.get('recommendation', 'No recommendation provided'),
            'image_path': finding.get('image_path', None)
        }
        formatted_findings.append(formatted_finding)

    # Load and render template
    template_loader = jinja2.FileSystemLoader('.')
    template_env = jinja2.Environment(loader=template_loader)
    template = template_env.get_template('campus-inspection-report.html')
    
    html_content = template.render(
        timestamp=timestamp,
        inspector_name=inspector_name,
        location=location,
        weather=weather,
        findings=formatted_findings
    )
    
    return html_content

@app.route('/health')
def health_check():
    try:
        # Simple test to verify API connection
        response = model.generate_content("Test connection")
        return jsonify({
            'status': 'healthy',
            'api_configured': True
        })
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}")
        return jsonify({
            'status': 'unhealthy',
            'error': str(e)
        }), 500

@app.route('/report_status/<task_id>')
def report_status(task_id):
    """Check the status of a report generation task"""
    try:
        # Implement status checking logic
        return jsonify({
            'status': 'processing',
            'progress': 50,  # Example progress percentage
            'message': 'Processing images...'
        })
    except Exception as e:
        logger.error(f"Error checking status: {str(e)}")
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)