"""Persist tab-based application drafts in `drafts.payload` (JSONB)."""

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Optional

from sqlalchemy import select
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.ext.asyncio import AsyncSession

from src.audit import AuditAction, record_audit, resolve_actor_fields
from src.initiative_db.application_storage import record_tab_snapshot
from src.initiative_db.models import AuditLog, Draft, Initiative

# Seed user from migrations/001_initiative_schema.sql
SYSTEM_DRAFT_OWNER_ID = uuid.UUID("00000000-0000-4000-8000-000000000001")
logger = logging.getLogger(__name__)
OFFICIAL_FORM_LAYOUT_PAYLOAD_KEY = "officialFormPdfLayout"


async def save_application_draft_tab(
    session: AsyncSession,
    case_id: str,
    tab: str,
    data: Dict[str, Any],
    owner_id: Optional[uuid.UUID] = None,
) -> Dict[str, Any]:
    stmt = select(Initiative).where(Initiative.case_code == case_id)
    initiative = (await session.execute(stmt)).scalar_one_or_none()

    effective_owner = owner_id or SYSTEM_DRAFT_OWNER_ID

    if initiative is None:
        initiative = Initiative(case_code=case_id, owner_id=effective_owner)
        session.add(initiative)
        await session.flush()
        draft = Draft(
            draft_code=f"DRAFT-{case_id}",
            initiative_id=initiative.id,
            payload=_empty_payload(case_id),
        )
        session.add(draft)
    else:
        stmt_d = (
            select(Draft)
            .where(Draft.initiative_id == initiative.id)
            .order_by(Draft.updated_at.desc())
            .limit(1)
        )
        draft = (await session.execute(stmt_d)).scalar_one_or_none()
        if draft is None:
            draft = Draft(
                draft_code=f"DRAFT-{case_id}",
                initiative_id=initiative.id,
                payload=_empty_payload(case_id),
            )
            session.add(draft)

    if owner_id and initiative.owner_id == SYSTEM_DRAFT_OWNER_ID:
        initiative.owner_id = owner_id

    current = dict(draft.payload) if isinstance(draft.payload, dict) else {}
    tabs = dict(current.get("tabs") or {})
    tabs[tab] = data
    now_iso = datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    current["caseId"] = case_id
    current["updatedAt"] = now_iso
    current["tabs"] = tabs

    draft.payload = current
    draft.version = (draft.version or 0) + 1

    await session.flush()

    if owner_id is not None:
        ae, ar = await resolve_actor_fields(session, owner_id)
        await record_audit(
            session,
            actor_user_id=owner_id,
            actor_email=ae,
            actor_role=ar,
            action=AuditAction.update,
            entity_type="application_draft",
            entity_id=case_id,
            after={"tab": tab, "version": draft.version},
            metadata={"source": "autosave_tab"},
        )

    # Snapshot history is optional telemetry.
    # Run it in a savepoint so missing table errors don't poison the main transaction.
    try:
        async with session.begin_nested():
            await record_tab_snapshot(
                session,
                initiative_id=initiative.id,
                draft_id=draft.id,
                tab=tab,
                payload=data,
                source="autosave",
            )
    except ProgrammingError as exc:
        if "draft_tab_snapshots" in str(exc).lower():
            logger.warning(
                "draft_tab_snapshots table missing; skipping tab snapshot for case %s (tab=%s)",
                case_id,
                tab,
            )
        else:
            raise

    session.add(
        AuditLog(
            actor_id=SYSTEM_DRAFT_OWNER_ID,
            action="update",
            entity="draft",
            entity_id=draft.id,
            diff={"tab": tab, "version": draft.version},
        )
    )

    return {
        "caseId": case_id,
        "updatedAt": now_iso,
        "tabs": tabs,
        "version": draft.version,
        "publicUrl": f"/application-drafts/{case_id}.yml",
    }


async def insert_initiative_draft_if_missing(
    session: AsyncSession, case_id: str, payload: Dict[str, Any]
) -> None:
    """
    Copy a YAML/file-shaped bundle into initiatives + drafts when Postgres was empty.
    Idempotent: no-op if initiative already exists.
    """
    stmt = select(Initiative).where(Initiative.case_code == case_id)
    if (await session.execute(stmt)).scalar_one_or_none() is not None:
        return

    merged: Dict[str, Any] = dict(payload) if isinstance(payload, dict) else {}
    tabs = merged.get("tabs")
    merged["tabs"] = dict(tabs) if isinstance(tabs, dict) else {}
    merged["caseId"] = case_id

    initiative = Initiative(case_code=case_id, owner_id=SYSTEM_DRAFT_OWNER_ID)
    session.add(initiative)
    await session.flush()
    draft = Draft(
        draft_code=f"DRAFT-{case_id}",
        initiative_id=initiative.id,
        payload=merged,
    )
    session.add(draft)
    await session.flush()


async def get_application_draft_document(session: AsyncSession, case_id: str) -> Dict[str, Any]:
    stmt = select(Initiative).where(Initiative.case_code == case_id)
    initiative = (await session.execute(stmt)).scalar_one_or_none()
    if initiative is None:
        raise KeyError(case_id)

    stmt_d = (
        select(Draft)
        .where(Draft.initiative_id == initiative.id)
        .order_by(Draft.updated_at.desc())
        .limit(1)
    )
    draft = (await session.execute(stmt_d)).scalar_one_or_none()
    if draft is None or not isinstance(draft.payload, dict):
        raise KeyError(case_id)
    return draft.payload


def _empty_payload(case_id: str) -> Dict[str, Any]:
    return {
        "caseId": case_id,
        "updatedAt": datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z"),
        "tabs": {},
    }


async def get_latest_draft_for_case(session: AsyncSession, case_id: str) -> Optional[Draft]:
    stmt = select(Initiative).where(Initiative.case_code == case_id)
    initiative = (await session.execute(stmt)).scalar_one_or_none()
    if initiative is None:
        return None
    stmt_d = (
        select(Draft)
        .where(Draft.initiative_id == initiative.id)
        .order_by(Draft.updated_at.desc())
        .limit(1)
    )
    return (await session.execute(stmt_d)).scalar_one_or_none()


async def get_official_form_layout_payload(session: AsyncSession, case_id: str) -> Optional[Dict[str, Any]]:
    draft = await get_latest_draft_for_case(session, case_id)
    if draft is None or not isinstance(draft.payload, dict):
        return None
    slot = draft.payload.get(OFFICIAL_FORM_LAYOUT_PAYLOAD_KEY)
    if not isinstance(slot, dict):
        return None
    return dict(slot)


async def save_official_form_layout_payload(
    session: AsyncSession,
    *,
    case_id: str,
    payload: Dict[str, Any],
    owner_id: Optional[uuid.UUID] = None,
) -> Dict[str, Any]:
    stmt = select(Initiative).where(Initiative.case_code == case_id)
    initiative = (await session.execute(stmt)).scalar_one_or_none()
    if initiative is None:
        initiative = Initiative(case_code=case_id, owner_id=owner_id or SYSTEM_DRAFT_OWNER_ID)
        session.add(initiative)
        await session.flush()

    stmt_d = (
        select(Draft)
        .where(Draft.initiative_id == initiative.id)
        .order_by(Draft.updated_at.desc())
        .limit(1)
    )
    draft = (await session.execute(stmt_d)).scalar_one_or_none()
    if draft is None:
        draft = Draft(
            draft_code=f"DRAFT-{case_id}",
            initiative_id=initiative.id,
            payload=_empty_payload(case_id),
        )
        session.add(draft)
        await session.flush()

    if owner_id and initiative.owner_id == SYSTEM_DRAFT_OWNER_ID:
        initiative.owner_id = owner_id

    now_iso = datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
    current = dict(draft.payload) if isinstance(draft.payload, dict) else _empty_payload(case_id)
    current["caseId"] = case_id
    current["updatedAt"] = now_iso
    current[OFFICIAL_FORM_LAYOUT_PAYLOAD_KEY] = dict(payload)
    draft.payload = current
    draft.version = (draft.version or 0) + 1
    await session.flush()
    return {"caseId": case_id, "updatedAt": now_iso, "version": draft.version}
