from __future__ import annotations

import copy

from django.core.management.base import BaseCommand

from wagtail import blocks
from wagtail.fields import StreamField
from wagtail.models import Page

from ocyan.plugin.wagtail.localization import localize_internal_path
from ocyan.plugin.wagtail.models import MenuLink


AI_SEARCH_PATHS = {
    "/ai-search/",
    "/smart-search/",
    "/wordspinner/ai-search/",
    "/wordspinner/smart-search/",
}


class Command(BaseCommand):
    help = (
        "Repair translated page chooser links and legacy AI Search URLs in "
        "localized Wagtail content."
    )

    def add_arguments(self, parser):
        parser.add_argument(
            "--source-language",
            default="nl",
            help="Canonical source language to exclude from repair.",
        )
        parser.add_argument(
            "--languages",
            nargs="*",
            help="Only repair these target languages.",
        )
        parser.add_argument(
            "--dry-run",
            action="store_true",
            help="Show what would change without saving.",
        )

    def handle(self, *args, **options):
        source_language = (options["source_language"] or "nl").strip().lower()
        target_languages = {
            lang.strip().lower().replace("_", "-")
            for lang in (options.get("languages") or [])
            if lang.strip()
        }
        dry_run = options["dry_run"]

        repaired_pages = 0
        repaired_menu_links = 0

        pages = Page.objects.exclude(locale__language_code=source_language).specific()
        if target_languages:
            pages = pages.filter(locale__language_code__in=sorted(target_languages))

        for page in pages:
            page_changed = False
            for field in page._meta.get_fields():
                if not isinstance(field, StreamField):
                    continue

                stream_value = getattr(page, field.name, None)
                raw_data = getattr(stream_value, "raw_data", None)
                if raw_data is None:
                    continue

                new_data, changed = self._repair_block_value(
                    field.stream_block,
                    copy.deepcopy(raw_data),
                    language_code=page.locale.language_code,
                )
                if changed:
                    setattr(page, field.name, new_data)
                    page_changed = True

            if not page_changed:
                continue

            repaired_pages += 1
            self.stdout.write(
                f"{'Would repair' if dry_run else 'Repaired'} page "
                f"{page.pk} [{page.locale.language_code}] {page.title}"
            )
            if dry_run:
                continue

            if page.live:
                page.save_revision().publish()
            else:
                page.save()

        menu_links = MenuLink.objects.exclude(locale__language_code=source_language)
        if target_languages:
            menu_links = menu_links.filter(locale__language_code__in=sorted(target_languages))

        for menu_link in menu_links:
            localized_link = localize_internal_path(
                menu_link.link,
                language_code=menu_link.locale.language_code,
            )
            if localized_link == menu_link.link:
                continue

            repaired_menu_links += 1
            self.stdout.write(
                f"{'Would repair' if dry_run else 'Repaired'} MenuLink "
                f"{menu_link.pk} [{menu_link.locale.language_code}] {menu_link.title}: "
                f"{menu_link.link} -> {localized_link}"
            )
            if dry_run:
                continue

            menu_link.link = localized_link
            if menu_link.live:
                menu_link.save_revision().publish()
            else:
                menu_link.save()

        self.stdout.write(
            self.style.SUCCESS(
                f"repair_localized_page_links complete: pages={repaired_pages}, "
                f"menu_links={repaired_menu_links}, dry_run={dry_run}"
            )
        )

    def _repair_block_value(self, block, value, language_code: str):
        changed = False

        if isinstance(block, blocks.StreamBlock):
            for child in value or []:
                child_block = block.child_blocks.get(child.get("type"))
                if child_block is None:
                    continue
                new_value, child_changed = self._repair_block_value(
                    child_block,
                    child.get("value"),
                    language_code=language_code,
                )
                if child_changed:
                    child["value"] = new_value
                    changed = True
            return value, changed

        if isinstance(block, blocks.StructBlock):
            for name, child_block in block.child_blocks.items():
                if not isinstance(value, dict) or name not in value:
                    continue
                new_value, child_changed = self._repair_block_value(
                    child_block,
                    value.get(name),
                    language_code=language_code,
                )
                if child_changed:
                    value[name] = new_value
                    changed = True
            return value, changed

        if isinstance(block, blocks.ListBlock):
            if not isinstance(value, list):
                return value, changed
            for idx, item in enumerate(value):
                wrapped = isinstance(item, dict) and "value" in item
                item_value = item.get("value") if wrapped else item
                new_value, child_changed = self._repair_block_value(
                    block.child_block,
                    item_value,
                    language_code=language_code,
                )
                if child_changed:
                    if wrapped:
                        item["value"] = new_value
                    else:
                        value[idx] = new_value
                    changed = True
            return value, changed

        if isinstance(block, blocks.PageChooserBlock) and isinstance(value, int):
            translated_page_id = self._resolve_page_id(value, language_code)
            if translated_page_id and translated_page_id != value:
                return translated_page_id, True
            return value, False

        if isinstance(value, str) and value in AI_SEARCH_PATHS:
            localized = localize_internal_path(value, language_code=language_code)
            if localized != value:
                return localized, True
            return value, False

        return value, changed

    def _resolve_page_id(self, page_id: int, language_code: str):
        page = Page.objects.filter(pk=page_id).first()
        if page is None:
            return page_id
        if page.locale.language_code == language_code:
            return page_id

        translation_key = getattr(page, "translation_key", None)
        if not translation_key:
            return page_id

        translated_page = (
            Page.objects.filter(
                translation_key=translation_key,
                locale__language_code=language_code,
            )
            .only("id")
            .first()
        )
        return translated_page.id if translated_page else page_id
