run_pipeline_test.py

"""
run_pipeline_test.py
---------------------
End-to-end test for the full Phase 1+2 pipeline:

  Phase 1 (auto-run):
    Idea Agent -> Character Agent -> Story Agent -> [PAUSE]

  Phase 2 (after approval):
    Screenplay Agent -> Dialogue Agent -> END

Usage:
    python run_pipeline_test.py
    python run_pipeline_test.py "Your idea here"

Prints:
  - Story concept, characters, 3-act structure  (after Phase 1)
  - Full screenplay scenes with dialogue        (after Phase 2)
  - SQLite + ChromaDB write verification
"""

import sys
import textwrap
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

SEP  = "-" * 65
SEP2 = "=" * 65


def _wrap(text: str, indent: int = 4, width: int = 80) -> str:
    prefix = " " * indent
    return textwrap.fill(text or "(empty)", width=width,
                         initial_indent=prefix, subsequent_indent=prefix)


# ─────────────────────────────────────────────────────────────────────────────
def _print_phase1(state: dict) -> None:
    concept    = state.get("concept")
    characters = state.get("characters", [])
    acts       = state.get("acts", [])

    print(f"\n{SEP2}")
    print("  STORY CONCEPT")
    print(SEP2)
    if concept:
        print(f"  Title          : {concept.title}")
        print(f"  Genre          : {concept.genre}")
        print(f"  Target Audience: {concept.target_audience}")
        print(f"  Mood           : {concept.mood}")
        print(f"\n  Theme:")
        print(_wrap(concept.theme))
        print(f"\n  Synopsis:")
        print(_wrap(concept.synopsis))

    print(f"\n{SEP2}")
    print(f"  CHARACTERS  ({len(characters)} generated)")
    print(SEP2)
    for i, char in enumerate(characters, 1):
        print(f"\n  [{i}] {char.name}  [{char.role.upper()}]")
        print(f"  {SEP}")
        print(f"  Personality  : ", end=""); print(_wrap(char.personality, indent=0))
        print(f"  Goals        : ", end=""); print(_wrap(char.goals, indent=0))
        print(f"  Weaknesses   : ", end=""); print(_wrap(char.weaknesses, indent=0))
        print(f"  Speaking Style: ", end=""); print(_wrap(char.speaking_style, indent=0))
        print(f"  Background   : ", end=""); print(_wrap(char.background, indent=0))
        if char.relationships:
            print(f"  Relationships:")
            for name, rel in char.relationships.items():
                print(f"    -> {name}: {rel}")

    print(f"\n{SEP2}")
    print(f"  STORY STRUCTURE  ({len(acts)} acts)")
    print(SEP2)
    for act in acts:
        print(f"\n  ACT {act.act_number}: {act.title}")
        print(f"  {SEP}")
        print(f"  Summary:")
        print(_wrap(act.summary))
        print(f"\n  Scenes:")
        for j, scene_summary in enumerate(act.scene_summaries, 1):
            global_num = (act.act_number - 1) * 5 + j
            print(f"\n    Scene {global_num}:")
            print(_wrap(scene_summary, indent=6))


def _print_phase2(state: dict) -> None:
    scenes    = state.get("scenes", [])
    dialogue  = state.get("dialogue", [])

    # Build dialogue lookup: scene_number -> list of lines in sequence order
    dlg_by_scene: dict[int, list] = {}
    for d in dialogue:
        dlg_by_scene.setdefault(d.scene_number, []).append(d)
    for lines in dlg_by_scene.values():
        lines.sort(key=lambda x: x.sequence)

    # Print first 3 scenes in full + summary count for the rest
    FULL_PRINT = 3
    print(f"\n{SEP2}")
    print(f"  SCREENPLAY + DIALOGUE  ({len(scenes)} scenes, {len(dialogue)} total lines)")
    print(SEP2)

    for scene in scenes:
        sn = scene.scene_number
        is_full = sn <= FULL_PRINT

        print(f"\n  {'=' * 55}")
        print(f"  SCENE {sn}  (Act {scene.act_number})  [{scene.emotion_tone}]")
        print(f"  {scene.slugline}")
        print(f"  {'=' * 55}")

        if is_full:
            print()
            # Action lines -- wrap each paragraph separately
            for para in (scene.action_lines or "").split("\n"):
                if para.strip():
                    print(_wrap(para, indent=2))
            print()

            lines = dlg_by_scene.get(sn, [])
            if lines:
                for dl in lines:
                    print(f"  {dl.character_name.upper()}")
                    print(_wrap(f"({dl.emotion}) {dl.line}", indent=10))
                    print()
            else:
                print("  (no dialogue generated for this scene)")
        else:
            print(f"  Summary: {scene.summary}")
            lines = dlg_by_scene.get(sn, [])
            print(f"  Dialogue: {len(lines)} lines")

    if len(scenes) > FULL_PRINT:
        print(f"\n  [Scenes {FULL_PRINT + 1}-{len(scenes)} shown as summaries above]")


# ─────────────────────────────────────────────────────────────────────────────
def main():
    print(SEP2)
    print("  Scriptify AI -- Full Pipeline Test  (Phase 1 + Phase 2)")
    print(SEP2)

    # ── Validate Ollama ───────────────────────────────────────────────────────
    if not Path(".env").exists():
        print("\n  ERROR: .env not found. Copy .env.example to .env.")
        sys.exit(1)

    from backend.config import require_ollama, OLLAMA_MODEL
    try:
        require_ollama()
        print(f"\n  [OK] Ollama running  |  model: {OLLAMA_MODEL}")
    except EnvironmentError as e:
        print(f"\n  ERROR: {e}")
        sys.exit(1)

    # ── Idea ──────────────────────────────────────────────────────────────────
    raw_idea = (
        sys.argv[1] if len(sys.argv) > 1
        else "A detective solving mysterious murders in Mumbai"
    )
    print(f"\n  Idea: \"{raw_idea}\"")

    # ── Create project ────────────────────────────────────────────────────────
    from backend.db.database import create_tables, SessionLocal
    from backend.db import crud
    create_tables()
    db = SessionLocal()
    try:
        project = crud.create_project(db, title="[Pipeline Test -- pending]")
        project_id = project.id
        print(f"  [OK] Project created  (id={project_id})")
    finally:
        db.close()

    # ── PHASE 1: idea -> character -> story -> pause ───────────────────────────
    import time
    print(f"\n{SEP}")
    print("  PHASE 1: Idea / Character / Story agents running...")
    print(f"  (3 x 1 LLM call for idea+characters, 3 calls for story acts)")
    print(SEP)

    t0 = time.time()
    from backend.agents.graph import run_pipeline
    try:
        thread_id, state = run_pipeline(project_id=project_id, raw_idea=raw_idea)
    except Exception as e:
        print(f"\n  PIPELINE ERROR: {e}")
        import traceback; traceback.print_exc()
        sys.exit(1)

    print(f"\n  Phase 1 complete ({time.time() - t0:.1f}s)")
    print(f"  Thread ID: {thread_id}")

    _print_phase1(state)

    # ── Verify Phase 1 SQLite ─────────────────────────────────────────────────
    print(f"\n{SEP2}")
    print("  SQLITE VERIFICATION (Phase 1)")
    print(SEP2)
    db = SessionLocal()
    try:
        proj     = crud.get_project(db, project_id)
        db_chars = crud.list_characters(db, project_id)
        db_acts  = crud.list_acts(db, project_id)
        print(f"\n  Project: {proj.title}  [status={proj.status}]")
        print(f"  Characters in DB: {len(db_chars)}")
        for c in db_chars:
            print(f"    [{c.id}] {c.name}  ({c.role})")
        print(f"  Acts in DB: {len(db_acts)}")
        for a in db_acts:
            print(f"    [{a.id}] Act {a.act_number}: {a.title}")
    finally:
        db.close()

    # ── Verify ChromaDB ───────────────────────────────────────────────────────
    from backend.memory.chroma_client import get_collection, COLLECTION_CHARACTER_MEMORY
    col   = get_collection(COLLECTION_CHARACTER_MEMORY)
    metas = col.get(where={"project_id": {"$eq": project_id}}).get("metadatas", [])
    print(f"\n  ChromaDB character_memory: {len(metas)} entries")
    for m in metas:
        print(f"    {m.get('character_name')}  ({m.get('role')})")

    # ── PHASE 2: approve + screenplay + dialogue ───────────────────────────────
    print(f"\n{SEP}")
    print("  PHASE 2: Approving review gate -> Screenplay + Dialogue agents...")
    print(f"  ({sum(len(a.scene_summaries) for a in state.get('acts', []))} scenes x 2 calls each)")
    print(f"  Expected time: ~{sum(len(a.scene_summaries) for a in state.get('acts', []))*2}+ minutes on CPU")
    print(SEP)

    t1 = time.time()
    from backend.agents.graph import resume_pipeline
    try:
        final_state = resume_pipeline(thread_id)
    except Exception as e:
        print(f"\n  PHASE 2 ERROR: {e}")
        import traceback; traceback.print_exc()
        sys.exit(1)

    print(f"\n  Phase 2 complete ({time.time() - t1:.1f}s)")

    _print_phase2(final_state)

    # ── Verify Phase 2 SQLite ─────────────────────────────────────────────────
    print(f"\n{SEP2}")
    print("  SQLITE VERIFICATION (Phase 2)")
    print(SEP2)
    db = SessionLocal()
    try:
        db_scenes = crud.list_scenes(db, project_id)
        print(f"\n  Scenes in DB: {len(db_scenes)}")
        total_dlg = 0
        for s in db_scenes[:5]:  # show first 5
            dlg = crud.list_dialogue(db, s.id)
            total_dlg += len(dlg)
            print(f"    Scene {s.scene_number}: \"{s.slugline}\"  [{s.emotion_tone}]  "
                  f"-- {len(dlg)} dialogue lines")
        if len(db_scenes) > 5:
            for s in db_scenes[5:]:
                dlg = crud.list_dialogue(db, s.id)
                total_dlg += len(dlg)
            print(f"    ... ({len(db_scenes) - 5} more scenes)")
        print(f"\n  Total dialogue lines in DB: {total_dlg}")
    finally:
        db.close()

    # ── Verify ChromaDB Phase 2 ───────────────────────────────────────────────
    from backend.memory.chroma_client import (
        COLLECTION_SCENE_SUMMARIES, COLLECTION_DIALOGUE_HISTORY
    )
    sc_col = get_collection(COLLECTION_SCENE_SUMMARIES)
    dh_col = get_collection(COLLECTION_DIALOGUE_HISTORY)
    sc_ids = sc_col.get(where={"project_id": {"$eq": project_id}}).get("ids", [])
    dh_ids = dh_col.get(where={"project_id": {"$eq": project_id}}).get("ids", [])
    print(f"\n  ChromaDB scene_summaries   : {len(sc_ids)} entries")
    print(f"  ChromaDB dialogue_history  : {len(dh_ids)} entries")

    # ── Done ──────────────────────────────────────────────────────────────────
    print(f"\n{SEP2}")
    print("  PIPELINE COMPLETE")
    print(SEP2)
    print(f"\n  Project ID : {project_id}")
    print(f"  Thread ID  : {thread_id}")
    print(f"  Total time : {time.time() - t0:.1f}s")
    print()


if __name__ == "__main__":
    main()