import easyocr
import torch
import numpy as np
from compliance_rules import ComplianceRules

# Initialize OCR reader
reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())

# Initialize compliance rules
compliance_rules = ComplianceRules()

def extract_text_from_image(image):
    """Extract text from image using EasyOCR"""
    try:
        result = reader.readtext(np.array(image))
        return " ".join([text[1] for text in result])
    except Exception as e:
        print(f"Error in text extraction: {str(e)}")
        return "Error extracting text from image"

def check_compliance(text, regions=None):
    """Check text for compliance across selected regions"""
    all_rules = compliance_rules.get_all_rules()
    
    # Filter rules based on selected regions
    if regions is not None:
        rules = {region: rules for region, rules in all_rules.items() if regions.get(region, True)}
    else:
        rules = all_rules

    report = {
        "compliant": True,
        "violations": [],
        "warnings": [],
        "channel_risks": {
            "email": {"score": 0, "details": []},
            "social": {"score": 0, "details": []},
            "print": {"score": 0, "details": []}
        }
    }
    
    for region, region_rules in rules.items():
        # Check prohibited terms
        for term_info in region_rules["prohibited_terms"]:
            term = term_info["term"].lower()
            if term in text.lower() or any(var.lower() in text.lower() for var in term_info["variations"]):
                report["compliant"] = False
                violation = f"{region}: Prohibited term '{term}' found"
                report["violations"].append({
                    "region": region,
                    "type": "prohibited_term",
                    "term": term,
                    "severity": term_info["severity"]
                })
                
                # Update channel risks
                for channel in report["channel_risks"]:
                    risk_score = compliance_rules.calculate_risk_score([violation], [], region)
                    report["channel_risks"][channel]["score"] += risk_score
                    report["channel_risks"][channel]["details"].append(
                        f"Prohibited term '{term}' increases {channel} risk"
                    )
        
        # Check required disclaimers
        for disclaimer in region_rules["required_disclaimers"]:
            disclaimer_found = any(
                disc_text.lower() in text.lower() 
                for disc_text in disclaimer["text"]
            )
            if not disclaimer_found:
                warning = f"{region}: Missing {disclaimer['type']} disclaimer"
                report["warnings"].append({
                    "region": region,
                    "type": "missing_disclaimer",
                    "disclaimer_type": disclaimer["type"],
                    "severity": disclaimer["severity"]
                })
                
                # Update channel risks
                for channel in report["channel_risks"]:
                    risk_score = compliance_rules.calculate_risk_score([], [warning], region)
                    report["channel_risks"][channel]["score"] += risk_score
                    report["channel_risks"][channel]["details"].append(
                        f"Missing {disclaimer['type']} disclaimer affects {channel} risk"
                    )
    
    return report

def format_severity(severity):
    return f'<span class="severity-{severity}">{severity.upper()}</span>'

def generate_html_report(compliance_report):
    """Generate formatted HTML report"""
    html = '<div class="report-container">'
    
    # Overall Status
    status_class = "compliant" if compliance_report["compliant"] else "non-compliant"
    status_icon = "✅" if compliance_report["compliant"] else "❌"
    html += f'<div class="status {status_class}">{status_icon} Overall Status: {"Compliant" if compliance_report["compliant"] else "Non-Compliant"}</div>'
    
    # Violations
    if compliance_report["violations"]:
        html += '<div class="section">'
        html += '<div class="section-title">🚫 Violations Found:</div>'
        for violation in compliance_report["violations"]:
            html += f'<div class="item">• {violation["region"]}: {violation["type"]} - \'{violation["term"]}\' (Severity: {format_severity(violation["severity"])})</div>'
        html += '</div>'
    
    # Warnings
    if compliance_report["warnings"]:
        html += '<div class="section">'
        html += '<div class="section-title">⚠️ Warnings:</div>'
        for warning in compliance_report["warnings"]:
            html += f'<div class="item">• {warning["region"]}: {warning["disclaimer_type"]} (Severity: {format_severity(warning["severity"])})</div>'
        html += '</div>'
    
    # Channel Risk Assessment
    html += '<div class="section">'
    html += '<div class="section-title">📊 Channel Risk Assessment:</div>'
    
    for channel, risk_info in compliance_report["channel_risks"].items():
        score = risk_info["score"]
        risk_level = "low" if score < 3 else "medium" if score < 6 else "high"
        
        html += f'<div class="channel risk-{risk_level}">'
        html += f'<strong>{channel.capitalize()}</strong>: {risk_level.upper()} Risk (Score: {score})'
        
        if risk_info["details"]:
            html += '<div class="details">'
            for detail in risk_info["details"]:
                html += f'<div>• {detail}</div>'
            html += '</div>'
        html += '</div>'
    
    html += '</div></div>'
    return html

def analyze_ad_copy(image, regions=None):
    """Main function to analyze ad copy"""
    # Extract text from image
    text = extract_text_from_image(image)
    
    # Check compliance
    compliance_report = check_compliance(text, regions)
    
    # Generate HTML report
    return generate_html_report(compliance_report)