from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
import json
from datetime import datetime, timezone

from .demo_pack_loader import DemoPack, MemorySeed, load_default_demo_pack
from .prompt_rewriter import RewriteResult, rewrite_memory
from .quilt import assemble_quilt, card_summaries, render_tile, save_image


@dataclass(frozen=True)
class QuiltState:
    cards: tuple[RewriteResult, ...] = ()
    style: str = "Fabric Quilt"
    pack_id: str = ""


@dataclass(frozen=True)
class RenderResult:
    quilt_path: Path
    tile_path: Path
    log_path: Path
    cards: tuple[RewriteResult, ...]


def empty_state(style: str = "Fabric Quilt") -> QuiltState:
    return QuiltState(cards=(), style=style, pack_id="")


def seed_state_from_pack(pack: DemoPack, count: int = 6) -> QuiltState:
    cards = tuple(
        rewrite_memory(memory.text, memory.location, pack.style)
        for memory in pack.memories[: min(count, len(pack.memories))]
    )
    return QuiltState(cards=cards, style=pack.style, pack_id=pack.pack_id)


def add_memory_to_state(state: QuiltState, memory_text: str, location_tag: str = "", style: str | None = None, photo_path: str | None = None) -> tuple[QuiltState, RewriteResult]:
    resolved_style = style or state.style or "Fabric Quilt"
    card = rewrite_memory(memory_text, location_tag, resolved_style, photo_path=photo_path)
    next_state = QuiltState(cards=state.cards + (card,), style=resolved_style, pack_id=state.pack_id)
    return next_state, card


def build_quilt_image(state: QuiltState):
    cards = list(state.cards)
    if not cards:
        raise ValueError("build_quilt_image requires at least one model-generated card")
    return assemble_quilt(cards)


def build_single_tile(card: RewriteResult):
    return render_tile(card)


def _timestamp_slug() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def render_artifacts(state: QuiltState, output_dir: str | Path, stem: str = "memory_quilt") -> RenderResult:
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    cards = list(state.cards)
    if not cards:
        raise ValueError("render_artifacts requires at least one model-generated card")
    quilt_image = assemble_quilt(cards)
    tile_source = render_tile(cards[-1])
    tile_backend = str(getattr(tile_source, "info", {}).get("backend", "flux-diffusers"))
    tile_image = tile_source.convert("RGBA")
    last_card = cards[-1]
    last_meta = last_card.inference_meta if isinstance(last_card.inference_meta, dict) else {}
    quilt_path = save_image(quilt_image, output_path / f"{stem}_quilt.png")
    tile_path = save_image(tile_image, output_path / f"{stem}_tile.png")
    log_path = output_path / f"{stem}_log.json"
    log_payload = {
        "rendered_at": datetime.now(timezone.utc).isoformat(),
        "stem": stem,
        "card_count": len(cards),
        "cards": card_summaries(cards),
        "tile_backend": tile_backend,
        "model_id": str(last_meta.get("model_id") or last_card.selected_model_id),
        "adapter_name": str(last_meta.get("adapter_name") or last_card.prompt_source),
        "checkpoint_path": str(last_meta.get("checkpoint_path") or last_card.checkpoint_path),
        "checkpoint_source": str(last_meta.get("checkpoint_source") or "unknown"),
        "generation_stats": last_meta.get("generation_stats", {}),
        "quilt_path": str(quilt_path),
        "tile_path": str(tile_path),
    }
    log_path.write_text(json.dumps(log_payload, indent=2, ensure_ascii=False), encoding="utf-8")
    return RenderResult(quilt_path=quilt_path, tile_path=tile_path, log_path=log_path, cards=tuple(cards))


def render_sample_pack(output_dir: str | Path, count: int = 6) -> tuple[QuiltState, RenderResult, DemoPack]:
    pack = load_default_demo_pack()
    state = seed_state_from_pack(pack, count=count)
    result = render_artifacts(state, output_dir, stem=f"{pack.pack_id}_sample")
    return state, result, pack


def state_summary(state: QuiltState) -> str:
    if not state.cards:
        return "0 stitched contributions"
    return f"{len(state.cards)} stitched contribution(s) using {state.style}"