routes.py

from flask import Blueprint, render_template, redirect, url_for, request, flash, send_file, jsonify
from flask_login import login_user, logout_user, login_required, current_user
from werkzeug.security import check_password_hash, generate_password_hash
from models import db, User, Organization, KRIDefinition, KRIScore, AuditLog
from sqlalchemy import func, and_
from datetime import datetime
import csv
import io
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
from reportlab.lib.utils import simpleSplit
import json
bp = Blueprint('routes', __name__)

def log_audit(action, details):
    user_id = current_user.id if current_user.is_authenticated else None
    log = AuditLog(user_id=user_id, action=action, details=details)
    db.session.add(log)
    db.session.commit()

def compute_overall_score(org_id, year):
    scores = KRIScore.query.join(KRIDefinition).filter(KRIScore.organization_id == org_id, KRIScore.year == year).all()
    if not scores: return None
    total = sum(s.contribution for s in scores)
    return round(total, 2)

def risk_level_from_score(score):
    if score <= 25: return "Low"
    elif score <= 50: return "Medium"
    elif score <= 75: return "High"
    else: return "Critical"

# ==================== AUTH ====================
@bp.route('/')
def index(): return redirect(url_for('routes.login'))

@bp.route('/login', methods=['GET', 'POST'])
def login():
    if current_user.is_authenticated:
        return redirect(url_for('routes.dashboard'))
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        user = User.query.filter_by(username=username).first()
        if user and check_password_hash(user.password_hash, password):
            login_user(user)
            log_audit("Login", f"User {username} logged in")
            return redirect(url_for('routes.dashboard'))
        flash('Invalid credentials')
    return render_template('login.html')

@bp.route('/logout')
@login_required
def logout():
    log_audit("Logout", f"User {current_user.username} logged out")
    logout_user()
    return redirect(url_for('routes.login'))

# ==================== DASHBOARD ====================
@bp.route('/dashboard')
@login_required
def dashboard():
    year = request.args.get('year', type=int, default=datetime.now().year)
    if current_user.role == 'admin':
        orgs = Organization.query.all()
        selected_org = request.args.get('org_id', type=int)
        if selected_org: org = Organization.query.get(selected_org)
        else: org = orgs[0] if orgs else None
    else:
        org = current_user.organization
        orgs = [org]
        selected_org = org.id

    if not org:
        flash("No organization assigned")
        return redirect(url_for('routes.logout'))

    available_years = db.session.query(KRIScore.year).filter_by(organization_id=org.id).distinct().order_by(KRIScore.year).all()
    available_years = [y[0] for y in available_years]
    if not available_years: available_years = [datetime.now().year]
    if available_years: next_year = max(available_years) + 1
    else: next_year = datetime.now().year + 1

    overall = compute_overall_score(org.id, year)
    risk_level = risk_level_from_score(overall) if overall is not None else "N/A"
    kri_scores = KRIScore.query.join(KRIDefinition).filter(
        KRIScore.organization_id == org.id, KRIScore.year == year
    ).add_columns(KRIDefinition.kri_name, KRIDefinition.weight_pct, KRIDefinition.thresholds).all()

    kri_list = []
    for score, kri_name, weight_pct, thresholds_json in kri_scores:
        kri_list.append({'name': kri_name, 'raw': score.raw_value, 'score': score.numerical_score, 'weight': weight_pct, 'contribution': score.contribution})
    return render_template('dashboard.html', org=org, year=year, overall=overall, risk_level=risk_level, kri_list=kri_list, orgs=orgs, selected_org=selected_org, role=current_user.role, available_years=available_years, next_year=next_year)

# ==================== ADD/EDIT KRI SCORES ====================
@bp.route('/add_scores', methods=['GET', 'POST'])
@login_required
def add_scores():
    if current_user.role != 'admin' and current_user.role != 'org_user':
        flash("Unauthorized")
        return redirect(url_for('routes.dashboard'))

    org_id = current_user.organization_id if current_user.role == 'org_user' else request.args.get('org_id', type=int)
    if not org_id:
        flash("Organization required")
        return redirect(url_for('routes.dashboard'))

    org = Organization.query.get(org_id)
    year = request.args.get('year', type=int, default=datetime.now().year)

    if request.method == 'POST':
        kris = KRIDefinition.query.all()
        for kri in kris:
            raw_val = request.form.get(f'kri_{kri.id}', type=float)
            if raw_val is not None:
                thresholds = json.loads(kri.thresholds)
                num_score = 100
                for t in thresholds:
                    if t["min"] <= raw_val <= t["max"]:
                        num_score = t["score"]
                        break
                contrib = num_score * (kri.weight_pct / 100.0)
                existing = KRIScore.query.filter_by(organization_id=org.id, year=year, kri_id=kri.id).first()
                if existing:
                    existing.raw_value = raw_val
                    existing.numerical_score = num_score
                    existing.contribution = contrib
                    existing.calculated_at = datetime.utcnow()
                else:
                    new_score = KRIScore(organization_id=org.id, year=year, kri_id=kri.id, raw_value=raw_val, numerical_score=num_score, contribution=contrib)
                    db.session.add(new_score)
        db.session.commit()
        log_audit("Update Scores", f"Updated KRI scores for {org.name} year {year}")
        flash("Scores saved successfully")
        return redirect(url_for('routes.dashboard', org_id=org.id, year=year))

    kris = KRIDefinition.query.all()
    existing_scores = {score.kri_id: score for score in KRIScore.query.filter_by(organization_id=org.id, year=year).all()}
    return render_template('add_scores.html', org=org, year=year, kris=kris, existing=existing_scores)

# ==================== TRENDS (Chart.js) ====================
@bp.route('/trends')
@login_required
def trends():
    org_id = request.args.get('org_id', type=int)
    if current_user.role == 'admin':
        orgs = Organization.query.all()
        if org_id: org = Organization.query.get(org_id)
        else: org = orgs[0] if orgs else None
    else:
        orgs = []
        org = current_user.organization
        if not org:
            flash("No organization assigned to your account.")
            return redirect(url_for('routes.dashboard'))

    if not org:
        flash("Organization not found")
        return redirect(url_for('routes.dashboard'))

    years = db.session.query(KRIScore.year).filter_by(organization_id=org.id).distinct().order_by(KRIScore.year).all()
    years = [y[0] for y in years]
    if not years:
        flash("No score data available for this organization.")
        return render_template('trends.html', org=org, orgs=orgs if current_user.role == 'admin' else None, year_scores=[], available_years=[])

    year_scores = []
    for y in years:
        score = compute_overall_score(org.id, y)
        risk_level = risk_level_from_score(score) if score is not None else "N/A"
        year_scores.append({ 'year': y, 'score': score if score is not None else 0, 'risk_level': risk_level })
    return render_template('trends.html', org=org, orgs=orgs if current_user.role == 'admin' else None, year_scores=year_scores, available_years=years)

@bp.route('/api/trends_data')
@login_required
def trends_data():
    org_id = request.args.get('org_id', type=int)
    if current_user.role != 'admin' and org_id != current_user.organization_id:
        org_id = current_user.organization_id
    org = Organization.query.get(org_id)
    if not org: return jsonify([])
    years = db.session.query(KRIScore.year).filter_by(organization_id=org.id).distinct().order_by(KRIScore.year).all()
    result = []
    for y in years: result.append({"year": y[0], "score": compute_overall_score(org.id, y[0])})
    return jsonify(result)

# ==================== COMPARISON ENGINE ====================
@bp.route('/compare', methods=['GET', 'POST'])
@login_required
def compare():
    orgs = Organization.query.all()
    available_years = db.session.query(KRIScore.year).distinct().order_by(KRIScore.year).all()
    available_years = [y[0] for y in available_years]
    if not available_years: available_years = [datetime.now().year]
    mode = request.form.get('mode') if request.method == 'POST' else None

    context = {
        'orgs': orgs,
        'mode': mode,
        'available_years': available_years,
        'company_chart_data': None,
        'company_table_data': None,
        'company_org_names': [],
        'selected_comp_orgs': [],
        'comp_year': None,
        'year_chart_data': None,
        'year_table_data': None,
        'year_org_id': None,
        'year_org_name': None,
        'selected_years': [],
        'is_admin': current_user.role == 'admin'
    }

    if request.method == 'POST':
        if mode == 'company':
            if current_user.role != 'admin':
                flash("Unauthorized to compare multiple companies.")
                return redirect(url_for('routes.dashboard'))

            org_ids = []
            for key in ['comp_org1', 'comp_org2', 'comp_org3']:
                val = request.form.get(key)
                if val and val != '': org_ids.append(int(val))
            org_ids = list(dict.fromkeys(org_ids))
            if len(org_ids) < 2:
                flash("Please select at least two different companies.")
                return render_template('compare.html', **context)

            comp_year = int(request.form.get('comp_year', datetime.now().year))
            selected_orgs = Organization.query.filter(Organization.id.in_(org_ids)).all()
            org_names = [org.name for org in selected_orgs]

            scores = []
            for org in selected_orgs:
                score = compute_overall_score(org.id, comp_year)
                scores.append(score if score else 0)
            company_chart_data = {'labels': org_names, 'data': scores}
            kris = KRIDefinition.query.all()
            table_rows = []
            for kri in kris:
                row = {'kri_name': kri.kri_name}
                for org in selected_orgs:
                    kri_score = KRIScore.query.filter_by( organization_id=org.id, year=comp_year, kri_id=kri.id).first()
                    row[org.name] = {
                        'raw': kri_score.raw_value if kri_score else 'N/A',
                        'score': kri_score.numerical_score if kri_score else 'N/A'
                    }
                table_rows.append(row)

            context.update({
                'mode': 'company',
                'company_chart_data': company_chart_data,
                'company_table_data': table_rows,
                'company_org_names': org_names,
                'selected_comp_orgs': org_ids,
                'comp_year': comp_year
            })
            log_audit("Comparison", f"Compared companies: {', '.join(org_names)} for year {comp_year}")

        elif mode == 'year':
            if current_user.role == 'admin': org_id = request.form.get('year_org_id', type=int)
            else:
                org_id = current_user.organization_id
                if not org_id:
                    flash("No organization assigned to your account.")
                    return redirect(url_for('routes.dashboard'))

            if not org_id:
                flash("Select an organization.")
                return render_template('compare.html', **context)

            years = []
            for key in ['year1', 'year2', 'year3']:
                val = request.form.get(key)
                if val and val != '': years.append(int(val))
            years = sorted(list(dict.fromkeys(years)))
            if len(years) < 2:
                flash("Please select at least two different years.")
                return render_template('compare.html', **context)

            org = Organization.query.get(org_id)
            if not org:
                flash("Organization not found.")
                return render_template('compare.html', **context)

            scores = []
            for y in years:
                score = compute_overall_score(org.id, y)
                scores.append(score if score else 0)

            year_chart_data = {
                'labels': [str(y) for y in years],
                'data': scores
            }

            kris = KRIDefinition.query.all()
            table_rows = []
            for kri in kris:
                row = {'kri_name': kri.kri_name}
                for y in years:
                    kri_score = KRIScore.query.filter_by( organization_id=org.id, year=y, kri_id=kri.id).first()
                    row[y] = {
                        'raw': kri_score.raw_value if kri_score else 'N/A',
                        'score': kri_score.numerical_score if kri_score else 'N/A'
                    }
                table_rows.append(row)

            context.update({
                'mode': 'year',
                'year_chart_data': year_chart_data,
                'year_table_data': table_rows,
                'year_org_id': org_id,
                'year_org_name': org.name,
                'selected_years': years
            })
            log_audit("Comparison", f"Compared years {', '.join(str(y) for y in years)} for organization {org.name}")
    return render_template('compare.html', **context)

# ==================== REPORT EXPORT (PDF/CSV) ====================
@bp.route('/report/<int:org_id>/<int:year>/<format>')
@login_required
def report(org_id, year, format):
    if current_user.role != 'admin' and current_user.organization_id != org_id:
        flash("Unauthorized")
        return redirect(url_for('routes.dashboard'))

    org = Organization.query.get(org_id)
    if not org:
        flash("Organization not found")
        return redirect(url_for('routes.dashboard'))

    overall = compute_overall_score(org.id, year)
    risk_level = risk_level_from_score(overall) if overall else "N/A"
    kri_scores = KRIScore.query.join(KRIDefinition).filter(
        KRIScore.organization_id == org.id, KRIScore.year == year
    ).add_columns(KRIDefinition.kri_name, KRIDefinition.weight_pct).all()

    if format == 'csv':
        output = io.StringIO()
        writer = csv.writer(output)
        writer.writerow(['KRI Name', 'Raw Value', 'Numerical Score', 'Weight (%)', 'Contribution'])
        for score, name, weight in kri_scores: writer.writerow([name, score.raw_value, score.numerical_score, weight, round(score.contribution,2)])
        writer.writerow([])
        writer.writerow(['Overall Score', overall, 'Risk Level', risk_level])
        output.seek(0)
        log_audit("Download Report", f"Downloaded CSV report for {org.name}, year {year}")
        return send_file(io.BytesIO(output.getvalue().encode('utf-8')), mimetype='text/csv', as_attachment=True, download_name=f"{org.name}_{year}_report.csv")

    elif format == 'pdf':
        from reportlab.lib.pagesizes import A4, portrait
        from reportlab.lib import colors
        from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer, PageBreak
        from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
        from reportlab.lib.units import inch, mm
        from reportlab.pdfgen import canvas
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(buffer, pagesize=A4, leftMargin=15*mm, rightMargin=15*mm, topMargin=25*mm, bottomMargin=20*mm)
        styles = getSampleStyleSheet()
        title_style = ParagraphStyle('CustomTitle', parent=styles['Heading1'], fontSize=16, spaceAfter=20, alignment=1)
        normal_style = styles['Normal']
        def add_header_footer(canvas_obj, doc):
            canvas_obj.saveState()
            canvas_obj.setFont('Helvetica-Bold', 10)
            header_text = "PTA RiskWatch - Compliance Monitoring Report"
            canvas_obj.drawString(15*mm, A4[1] - 15*mm, header_text)
            canvas_obj.line(15*mm, A4[1] - 18*mm, A4[0] - 15*mm, A4[1] - 18*mm)
            canvas_obj.setFont('Courier', 10)
            page_num = canvas_obj.getPageNumber()
            footer_text = f"Page {page_num}"
            canvas_obj.drawRightString(A4[0] - 15*mm, 15*mm, footer_text)
            canvas_obj.restoreState()
        story = []
        story.append(Paragraph(f"Compliance Report: {org.name}", title_style))
        story.append(Paragraph(f"Reporting Year: {year}", styles['Heading1']))
        story.append(Spacer(1, 10))
        score_text = f"Overall Risk Score: {overall}  |  Risk Level: {risk_level}"
        story.append(Paragraph(score_text, styles['Heading2']))
        story.append(Spacer(1, 0.2*inch))
        table_data = [['KRI Name', 'Raw (%)', 'Score', 'Wt. (%)', 'Contribution']]
        for score, name, weight in kri_scores:
            short_name = name if len(name) <= 45 else name[:42] + "..."
            table_data.append([
                short_name,
                f"{score.raw_value:.1f}" if isinstance(score.raw_value, float) else str(score.raw_value),
                str(score.numerical_score),
                str(weight),
                f"{score.contribution:.2f}"
            ])
        col_widths = [0.45*A4[0], 0.12*A4[0], 0.12*A4[0], 0.12*A4[0], 0.12*A4[0]]
        table = Table(table_data, colWidths=col_widths, repeatRows=1)
        table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#1e3c72')),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.white),
            ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
            ('ALIGN', (0, 1), (0, -1), 'LEFT'),
            ('FONTNAME', (0, 0), (-1, 0), 'Courier-Bold'),
            ('FONTSIZE', (0, 0), (-1, -1), 8),
            ('BOTTOMPADDING', (0, 0), (-1, 0), 8),
            ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
            ('GRID', (0, 0), (-1, -1), 0.5, colors.grey),
            ('LEFTPADDING', (0, 0), (-1, -1), 4),
            ('RIGHTPADDING', (0, 0), (-1, -1), 4),
            ('VALIGN', (0, 0), (-1, -1), 'MIDDLE'),
            ('TOPPADDING', (0, 1), (-1, -1), 3),
            ('BOTTOMPADDING', (0, 1), (-1, -1), 3),
        ]))
        story.append(table)
        story.append(Spacer(1, 0.2*inch))
        story.append(Paragraph(f"Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", normal_style))
        doc.build(story, onFirstPage=add_header_footer, onLaterPages=add_header_footer)
        buffer.seek(0)
        log_audit("Download Report", f"Downloaded PDF report for {org.name}, year {year}")
        return send_file(buffer, mimetype='application/pdf', as_attachment=True, download_name=f"{org.name}_{year}_report.pdf")

    else:
        flash("Invalid format")
        return redirect(url_for('routes.dashboard'))

# ==================== COMPANY REGISTRATION (PUBLIC) ====================
@bp.route('/register', methods=['GET', 'POST'])
def register():
    if current_user.is_authenticated: return redirect(url_for('routes.dashboard'))
    if request.method == 'POST':
        name = request.form.get('name')
        sector_type = request.form.get('sector_type')
        registration_date_str = request.form.get('registration_date')
        username = request.form.get('username')
        password = request.form.get('password')
        confirm_password = request.form.get('confirm_password')

        if not all([name, sector_type, registration_date_str, username, password]):
            flash("All fields are required.")
            return render_template('register.html')

        if password != confirm_password:
            flash("Passwords do not match.")
            return render_template('register.html')

        if User.query.filter_by(username=username).first():
            flash("Username already taken. Please choose another.")
            return render_template('register.html')

        if Organization.query.filter_by(name=name).first():
            flash("Organization with this name is already registered.")
            return render_template('register.html')

        try: registration_date = datetime.strptime(registration_date_str, '%Y-%m-%d').date()
        except ValueError:
            flash("Invalid date format. Use YYYY-MM-DD.")
            return render_template('register.html')

        new_org = Organization( name=name, sector_type=sector_type, registration_date=registration_date )
        db.session.add(new_org)
        db.session.flush()
        new_user = User( username=username, password_hash=generate_password_hash(password), role='org_user', organization_id=new_org.id )
        db.session.add(new_user)
        db.session.commit()

        log_audit("Company Registration", f"New company '{name}' registered by user '{username}'")
        flash("Registration successful! Please log in.")
        return redirect(url_for('routes.login'))
    return render_template('register.html')

# ==================== USER PROFILE MANAGEMENT ====================
@bp.route('/profile', methods=['GET', 'POST'])
@login_required
def profile():
    user = current_user
    org = user.organization if user.role == 'org_user' else None
    if request.method == 'POST':
        new_username = request.form.get('username', '').strip()
        new_password = request.form.get('new_password', '')
        confirm_password = request.form.get('confirm_password', '')

        if new_username and new_username != user.username:
            existing = User.query.filter_by(username=new_username).first()
            if existing:
                flash("Username already taken.")
                return render_template('profile.html', user=user, org=org)
            user.username = new_username
            log_audit("Profile Update", f"Changed username to '{new_username}'")

        if new_password:
            if new_password != confirm_password:
                flash("Passwords do not match.")
                return render_template('profile.html', user=user, org=org)
            if len(new_password) < 6:
                flash("Password must be at least 6 characters.")
                return render_template('profile.html', user=user, org=org)
            user.password_hash = generate_password_hash(new_password)
            log_audit("Profile Update", "Changed password")

        db.session.commit()
        flash("Profile updated successfully.")
        return redirect(url_for('routes.profile'))
    return render_template('profile.html', user=user, org=org)

# ==================== AUDIT LOG ====================
@bp.route('/audit')
@login_required
def audit():
    if current_user.role != 'admin':
        flash("Admin only")
        return redirect(url_for('routes.dashboard'))
    logs = AuditLog.query.order_by(AuditLog.timestamp.desc()).limit(200).all()
    return render_template('audit.html', logs=logs)