from __future__ import annotations

import importlib.util
import json
from pathlib import Path
from types import SimpleNamespace
import sys

ROOT = Path(__file__).resolve().parents[1]
for candidate in (ROOT, ROOT / "src"):
    if str(candidate) not in sys.path:
        sys.path.insert(0, str(candidate))

from apps.p5_memory_quilt.pipeline import QuiltState
from apps.p5_memory_quilt.prompt_rewriter import RewriteResult


def _load_headless_smoke_module():
    script_path = ROOT / "scripts" / "headless_smoke.py"
    spec = importlib.util.spec_from_file_location("p5_headless_smoke", script_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot import {script_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _dummy_card(model_id: str, checkpoint_path: str, *, location_tag: str = "corner store") -> RewriteResult:
    return RewriteResult(
        memory_text="Every Friday the tamale cart parked outside the blue house.",
        location_tag=location_tag,
        style="Fabric Quilt",
        preferred_model_id="jetbrains/Mellm-2-12B-Instruct",
        fallback_model_id="openbmb/MiniCPM5-1B",
        selected_model_id=model_id,
        checkpoint_path=checkpoint_path,
        prompt_source="local-transformers",
        caption="Corner Store",
        story="A stitched memory of the corner store.",
        flux_prompt="fabric quilt scene of a corner store at soft evening light",
        style_hint="stitched fabric texture",
        keywords=("tamale", "house"),
        season_hint="soft evening light and a calm neighborhood mood",
        inference_meta={
            "adapter_name": "local-transformers",
            "backend": "local-transformers",
            "model_id": model_id,
            "checkpoint_path": checkpoint_path,
            "checkpoint_source": "primary-checkpoint",
            "generation_stats": {
                "prompt_tokens": 11,
                "generated_tokens": 22,
                "elapsed_ms": 12.5,
            },
        },
        photo_path=None,
    )


def test_headless_smoke_writes_trace_artifact(tmp_path: Path, monkeypatch) -> None:
    module = _load_headless_smoke_module()
    module.TRACE_ARTIFACT_ROOT = tmp_path / "trace_artifacts"
    pack_path = ROOT / "data" / "demo_packs" / "p5_memory_quilt"
    output_dir = tmp_path / "verification"

    checkpoint = tmp_path / "checkpoint"
    checkpoint.mkdir()

    seed_card = _dummy_card("jetbrains/Mellm-2-12B-Instruct", str(checkpoint))
    smoke_card = _dummy_card("jetbrains/Mellm-2-12B-Instruct", str(checkpoint))

    def fake_seed_state_from_pack(pack, count: int = 6):
        return QuiltState(cards=(seed_card,), style=pack.style, pack_id=pack.pack_id)

    def fake_add_memory_to_state(current_state, memory_text: str, location_tag: str = "", style: str | None = None, photo_path: str | None = None):
        next_card = smoke_card
        next_state = QuiltState(cards=current_state.cards + (next_card,), style=style or current_state.style, pack_id=current_state.pack_id)
        return next_state, next_card

    def fake_render_artifacts(state: QuiltState, output_dir: Path, stem: str = "memory_quilt"):
        output_dir.mkdir(parents=True, exist_ok=True)
        quilt_path = output_dir / f"{stem}_quilt.png"
        tile_path = output_dir / f"{stem}_tile.png"
        log_path = output_dir / f"{stem}_log.json"
        quilt_path.write_bytes(b"quilt")
        tile_path.write_bytes(b"tile")
        log_path.write_text(json.dumps({"cards": len(state.cards)}), encoding="utf-8")
        return SimpleNamespace(quilt_path=quilt_path, tile_path=tile_path, log_path=log_path, cards=state.cards)

    monkeypatch.setattr(module, "seed_state_from_pack", fake_seed_state_from_pack)
    monkeypatch.setattr(module, "add_memory_to_state", fake_add_memory_to_state)
    monkeypatch.setattr(module, "render_artifacts", fake_render_artifacts)

    payload = module.run_smoke(pack_path, output_dir, count=6)
    trace_path = Path(payload["trace_path"])

    assert payload["repo"] == "all4-p5-memory-quilt"
    assert payload["project"] == "p5"
    assert payload["kind"] == "smoke"
    assert payload["network"] == "blocked"
    assert payload["model_name"] == "jetbrains/Mellm-2-12B-Instruct"
    assert payload["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
    assert payload["adapter_name"] == "local-transformers"
    assert payload["generation_stats"]["generated_tokens"] == 22
    assert trace_path.exists()

    trace = json.loads(trace_path.read_text(encoding="utf-8"))
    assert trace["kind"] == "smoke"
    assert trace["project"] == "p5"
    assert trace["pack_id"] == "p5_memory_quilt"
    assert trace["model_name"] == "jetbrains/Mellm-2-12B-Instruct"
    assert trace["model_id"] == "jetbrains/Mellm-2-12B-Instruct"
    assert trace["adapter_name"] == "local-transformers"
    assert trace["generation_stats"]["prompt_tokens"] == 11
    assert trace["inputs"]["memory_count"] == trace["parsed_outputs"]["tile_count"]
    assert Path(trace["parsed_outputs"]["quilt_path"]).exists()
    assert Path(trace["parsed_outputs"]["tile_path"]).exists()
    assert payload["trace_path"] == str(trace_path)