from django.contrib import admin
from .models import Document, DocumentSegment, Organization, Risk, Control, DocumentTemplate, DocumentRiskControl, DemoCode
from django.urls import reverse, path
from django.utils.html import format_html
from .utils import generate_demo_code, get_controls_for_risk, generate_key_findings, generate_recommendations
from .tables import get_risk_table
from django.shortcuts import render, redirect
from .forms import GenerateCodesForm
from django.conf import settings
from backend.accounts.utils import send_document_email
from django import forms
from django.contrib.admin.widgets import FilteredSelectMultiple

class DocumentRiskControlInline(admin.TabularInline):
    model = DocumentRiskControl
    extra = 2
    max_num = 10
    can_delete = False
    fields = ('risk', 'control', 'weight', 'likelihood')

    def get_formset(self, request, obj=None, **kwargs):
        formset = super().get_formset(request, obj, **kwargs)
        try:
            if request.method == 'POST' and 'organization_risks' in request.POST:
                risk_ids = request.POST.getlist('organization_risks')
                formset.form.base_fields['risk'].queryset = Risk.objects.filter(pk__in=risk_ids)
            elif obj:
                formset.form.base_fields['risk'].queryset = obj.organization.risks.all()
        except Exception:
            pass
        return formset


class DocumentAdminForm(forms.ModelForm):
    organization_risks = forms.ModelMultipleChoiceField(
        queryset=Risk.objects.all(),
        required=False,
        widget=FilteredSelectMultiple(verbose_name="Risks", is_stacked=False),
        help_text="Edit the AI-selected risks for this organization."
    )

    class Meta:
        model = Document
        fields = ['organization', 'status', 'key_findings', 'recomendations']

    class Media:
        css = { 'all': ('admin/css/widgets.css',) }
        js = (
            'admin/js/SelectBox.js',
            'admin/js/SelectFilter2.js',
        )

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if self.instance and getattr(self.instance, 'organization_id', None):
            self.fields['organization_risks'].initial = self.instance.organization.risks.all()


class DocumentAdmin(admin.ModelAdmin):
    change_form_template = "admin/core/document/change_form.html"

    form = DocumentAdminForm
    inlines = [DocumentRiskControlInline]
    list_display = ('organization', 'status', 'created_at', 'modified_at', 'review_link')
    list_filter = ('status', 'created_at')
    search_fields = ['organization__name', 'organization__email']
    readonly_fields = (
        'created_at', 'modified_at',
        'regen_controls_action', 'regen_keyfindings_action', 'regen_recommendations_action',
    )

    fieldsets = (
        ('Organization & Risks', {
            'fields': ('organization', 'organization_risks', 'regen_controls_action')
        }),
        ('Key Findings', {
            'fields': ('key_findings', 'regen_keyfindings_action')
        }),
        ('Recommendations', {
            'fields': ('recomendations', 'regen_recommendations_action')
        }),
        ('Status', {
            'fields': ('status',)
        }),
        ('Timestamps', {
            'fields': ('created_at', 'modified_at')
        }),
    )

    def regen_controls_action(self, obj):
        return format_html('<button type="submit" name="_regen_controls" class="button">Regenerate Controls using AI</button>')
    regen_controls_action.short_description = ''

    def regen_keyfindings_action(self, obj):
        return format_html('<button type="submit" name="_regen_key_findings" class="button">Regenerate Key Findings using AI</button>')
    regen_keyfindings_action.short_description = ''

    def regen_recommendations_action(self, obj):
        return format_html('<button type="submit" name="_regen_recommendations" class="button">Regenerate Recommendations using AI</button>')
    regen_recommendations_action.short_description = ''

    def save_model(self, request, obj, form, change):
        super().save_model(request, obj, form, change)
        org_risks = form.cleaned_data.get('organization_risks')
        if org_risks is not None and obj.organization_id:
            obj.organization.risks.set(org_risks)

    def _apply_post_org_risks(self, request, obj):
        try:
            if 'organization_risks' in request.POST and obj.organization_id:
                risk_ids = [int(pk) for pk in request.POST.getlist('organization_risks') if pk]
                obj.organization.risks.set(Risk.objects.filter(pk__in=risk_ids))
        except Exception:
            pass

    def _regen_controls(self, obj):
        obj.segments.filter(content__startswith="Identified Risks").delete()
        obj.segments.filter(content__startswith="Mitigation Controls").delete()
        obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()
        obj.documentriskcontrol_set.all().delete()
        top_risks = list(obj.organization.risks.all())
        obj.add_segment('h1', "Top 10 Risks Identified")
        risk_content = "\n\n".join([
            f"Risk: {r.risk_id} - {r.risk_name} \n"
            f"Category: {r.category}\n"
            f"Primary Impact: {r.primary_impact} \n"
            f"Secondary Impact: {r.secondary_impact}\n"
            f"Tertiary Impact: {r.tretiary_impact} \n"
            f"Detection Difficulty: {r.detection_difficulty} \n"
            f"Recovery Complexity: {r.recovery_complexity} \n"
            f"Business Impact Severity: {r.businnes_impact_severity}\n"
            for r in top_risks
        ])
        obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")
        controls_content = "Mitigation Controls:\n\n"
        for risk in top_risks:
            controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
            selected_controls = get_controls_for_risk(risk, organization=obj.organization)
            for control_id, weight, likelihood in selected_controls:
                control = Control.objects.filter(id=control_id).first()
                if control:
                    DocumentRiskControl.objects.create(
                        document=obj,
                        risk=risk,
                        control=control,
                        weight=weight,
                        likelihood=likelihood
                    )
                    label = f"{control.subcategory} - {control.function or ''}".rstrip(" -")
                    controls_content += f"  - Control: {label} (Impact Weight: {weight}/10) (Likelihood: {likelihood}/10)\n"
            controls_content += "\n"
        obj.add_segment('body', controls_content)

    def _regen_key_findings(self, obj):
        risks_top3 = get_risk_table(obj)[:3]
        key_findings = generate_key_findings(obj, risks_top3)
        if key_findings:
            obj.key_findings = key_findings
            obj.save(update_fields=['key_findings', 'modified_at'])
            return True
        return False

    def _regen_recommendations(self, obj):
        risks_top10 = get_risk_table(obj)[:10]
        recommendations = generate_recommendations(risks_top10, obj.organization)
        if recommendations:
            obj.recomendations = recommendations
            obj.save(update_fields=['recomendations', 'modified_at'])
            return True
        return False

    def changeform_view(self, request, object_id=None, form_url='', extra_context=None):
        if request.method == 'POST' and any(k in request.POST for k in ("_regen_controls", "_regen_key_findings", "_regen_recommendations")):
            obj = self.get_object(request, object_id)
            if obj is None:
                return super().changeform_view(request, object_id, form_url, extra_context)
            try:
                self._apply_post_org_risks(request, obj)
                if "_regen_controls" in request.POST:
                    self._regen_controls(obj)
                    self.message_user(request, "Risks and controls regenerated successfully.")
                elif "_regen_key_findings" in request.POST:
                    if self._regen_key_findings(obj):
                        self.message_user(request, "Key Findings regenerated.")
                    else:
                        self.message_user(request, "Key Findings could not be generated.", level='warning')
                elif "_regen_recommendations" in request.POST:
                    if self._regen_recommendations(obj):
                        self.message_user(request, "Recommendations regenerated.")
                    else:
                        self.message_user(request, "Recommendations could not be generated.", level='warning')
            except Exception as e:
                self.message_user(request, f"Action failed: {e}", level='error')
            return redirect(reverse('admin:core_document_change', args=[obj.pk]))
        return super().changeform_view(request, object_id, form_url, extra_context)

    def review_link(self, obj):
        url = reverse('admin:core_document_change', args=[obj.pk])
        label = 'Review / Edit'
        return format_html('<a class="button" href="{}">{}</a>', url, label)
    review_link.short_description = 'Action'

    def response_change(self, request, obj):
        if "_save_send" in request.POST:
            try:
                url = f"{settings.SITE_DOMAIN}/pdf/{obj.id}/"
                send_document_email(obj.organization.email, url, obj)
                obj.status = Document.STATUS_DONE
                obj.save(update_fields=['status', 'modified_at'])
                self.message_user(request, "Document sent and marked as done.")
            except Exception as e:
                self.message_user(request, f"Failed to send document: {e}", level='error')
            return redirect(reverse('admin:core_document_change', args=[obj.pk]))
        return super().response_change(request, obj)

    def _refresh_segments_from_current_mappings(self, obj):

        obj.segments.filter(content__startswith="Identified Risks").delete()
        obj.segments.filter(content__startswith="Mitigation Controls").delete()
        obj.segments.filter(content__in=["Top 10 Risks Identified", "Regenerated Controls"]).delete()

        top_risks = list(obj.organization.risks.all())

        if top_risks:
            obj.add_segment('h1', "Top 10 Risks Identified")
            risk_content = "\n\n".join([
                f"Risk: {r.risk_id} - {r.risk_name} \n"
                f"Category: {r.category}\n"
                f"Primary Impact: {r.primary_impact} \n"
                f"Secondary Impact: {r.secondary_impact}\n"
                f"Tertiary Impact: {r.tretiary_impact} \n"
                f"Detection Difficulty: {r.detection_difficulty} \n"
                f"Recovery Complexity: {r.recovery_complexity} \n"
                f"Business Impact Severity: {r.businnes_impact_severity}\n"
                for r in top_risks
            ])
            obj.add_segment('body', f"Identified Risks: \n\n{risk_content}")

        from collections import defaultdict
        controls_by_risk = defaultdict(list)
        for drc in obj.documentriskcontrol_set.select_related('risk', 'control').all():
            controls_by_risk[drc.risk_id].append(drc)

        controls_content = "Mitigation Controls:\n\n"
        for risk in top_risks:
            controls_content += f"Risk: {risk.risk_id} - {risk.risk_name}\n"
            rows = controls_by_risk.get(risk.pk, [])
            rows.sort(key=lambda x: (x.control.subcategory or '', x.control.function or ''))
            for drc in rows:
                control = drc.control
                if control:
                    label = f"{control.subcategory} - {control.function or ''}".rstrip(" -")
                    controls_content += f"  - Control: {label} (Impact Weight: {drc.weight}/10) (Likelihood: {drc.likelihood}/10)\n"
            controls_content += "\n"
        obj.add_segment('body', controls_content)

    def save_related(self, request, form, formsets, change):
        super().save_related(request, form, formsets, change)
        obj = form.instance
        try:
            self._refresh_segments_from_current_mappings(obj)
        except Exception:
            pass

class DocumentTemplateAdmin(admin.ModelAdmin):
    list_display = ['name', 'created_at', 'updated_at', 'preview_button']

    def preview_button(self, obj):
        url = reverse('core:template_preview', args=[obj.name])
        return format_html('<a href="{}" target="_blank">Preview</a>', url)

    preview_button.short_description = 'Preview'
    preview_button.allow_tags = True


class OrganizationAdmin(admin.ModelAdmin):
    list_display = ('name', 'email', 'industry_sector')
    search_fields = ['name', 'email']

class RiskAdmin(admin.ModelAdmin):
    ordering = ['risk_id']
    list_display = ['risk_id','risk_name','category']

class ControlAdmin(admin.ModelAdmin):
    list_display = ('id', 'subcategory', 'function', 'category')
    search_fields = ('subcategory', 'function', 'category')

class DocumentRiskControlAdmin(admin.ModelAdmin):
    list_display = ('document', 'risk', 'control', 'weight','likelihood')

class DemoCodeAdmin(admin.ModelAdmin):

    list_display = ('code', 'created_at', 'used', 'company', 'used_at')
    change_list_template = "admin/democode_changelist.html"

    def get_urls(self):
        urls = super().get_urls()
        custom_urls = [
            path('generate-codes/', self.admin_site.admin_view(self.generate_codes_view), name='generate-codes'),
        ]
        return custom_urls + urls

    def generate_codes_view(self, request):
        if request.method == 'POST':
            form = GenerateCodesForm(request.POST)
            if form.is_valid():
                count = form.cleaned_data['count']
                created = 0
                for _ in range(count):
                    while True:
                        code = generate_demo_code()
                        if not DemoCode.objects.filter(code=code).exists():
                            DemoCode.objects.create(code=code)
                            created += 1
                            break
                self.message_user(request, f"{created} codes generated.")
                return redirect('..')
        else:
            form = GenerateCodesForm()
        return render(request, 'admin/generate_codes.html', {'form': form})
    

admin.site.register(Document, DocumentAdmin)
admin.site.register(Organization, OrganizationAdmin)
admin.site.register(Risk ,RiskAdmin)
admin.site.register(Control, ControlAdmin)
admin.site.register(DocumentTemplate, DocumentTemplateAdmin)
admin.site.register(DocumentRiskControl, DocumentRiskControlAdmin)
admin.site.register(DemoCode, DemoCodeAdmin)