#!/usr/bin/env python3
# Flyto2 Core - Source Available License
# Copyright (c) 2025 Flyto2. All Rights Reserved.
#
# This source code is licensed under the Flyto2 Source Available License v1.0.
# Commercial use requires a license. See LICENSE for details.

"""
flyto-core Quickstart — see a data pipeline with evidence + replay in 30 seconds.

Usage:
    pip install flyto-core[api]
    python -m core.quickstart
"""

import asyncio
import json
import os
import shutil
import sys
import time
from pathlib import Path

os.environ["FLYTO_VALIDATION_MODE"] = "dev"

# ── Helpers ──────────────────────────────────────────────────────────

G = "\033[32m"  # green
R = "\033[31m"  # red
Y = "\033[33m"  # yellow
C = "\033[36m"  # cyan
D = "\033[90m"  # dim
B = "\033[1m"   # bold
P = "\033[35m"  # purple
X = "\033[0m"   # reset

LOGO = f"""
  {P}┌─┐┬  ┬ ┬┌┬┐┌─┐{X}
  {P}├┤ │  └┬┘ │ │ │{X}  {B}core{X} {D}v2.4{X}
  {P}└  ┴─┘ ┴  ┴ └─┘{X}  {D}Deterministic execution for AI agents{X}
"""

SAMPLE_JSON = '{"team": "Alpha", "members": [{"name": "Alice", "role": "Engineer"}, {"name": "Bob", "role": "Designer"}, {"name": "Carol", "role": "Manager"}], "active": true}'


def header(title):
    w = 64
    print(f"\n{'─' * w}")
    print(f"  {B}{title}{X}")
    print(f"{'─' * w}")


def ok(msg):
    print(f"  {G}✓{X} {msg}")


def fail(msg):
    print(f"  {R}✗{X} {msg}")


def info(msg):
    print(f"  {D}{msg}{X}")


def summarize_output(data):
    """Extract a short display string from module output."""
    if "result" in data:
        return str(data["result"]).replace("\n", " ")[:50]
    if "row_count" in data:
        cols = data.get("columns", [])
        return f"{data['row_count']} rows × {len(cols)} cols ({', '.join(cols)})"
    if "data" in data and isinstance(data["data"], dict):
        keys = list(data["data"].keys())
        return f"parsed: {{ {', '.join(keys)} }}"
    if "content" in data:
        preview = str(data["content"])[:40].replace("\n", " ")
        return preview
    if "bytes_written" in data:
        p = data.get("path", "?")
        # Show only filename, not full path
        short = p.split("/")[-1] if "/" in str(p) else str(p)
        return f"{data['bytes_written']} bytes → {short}"
    for k, v in data.items():
        return f"{k}: {str(v)[:40]}"
    return "(no output)"


# ── Main ─────────────────────────────────────────────────────────────

async def main():
    print(LOGO)

    work_dir = Path("./_quickstart_data")
    evidence_path = Path("./_quickstart_evidence")

    try:
        await _run_demo(work_dir, evidence_path)
    finally:
        await asyncio.sleep(0.3)  # let async evidence writes flush
        for d in [evidence_path, work_dir]:
            if d.exists():
                shutil.rmtree(d, ignore_errors=True)


async def _run_demo(work_dir, evidence_path):
    from core.engine import WorkflowEngine
    from core.engine.evidence import create_evidence_store
    from core.api.evidence_hooks import APIEvidenceHooks

    evidence_store = create_evidence_store(evidence_path)

    # ── Prepare sample data ───────────────────────────────────────
    work_dir.mkdir(parents=True, exist_ok=True)
    json_path = str(work_dir / "team.json")
    report_path = str(work_dir / "report.txt")

    with open(json_path, "w") as f:
        f.write(SAMPLE_JSON)

    # ── 1. Run a 5-step data pipeline ─────────────────────────────

    header("1. Execute Data Pipeline (5 steps, 4 categories)")

    workflow = {
        "name": "data-pipeline-demo",
        "steps": [
            {
                "id": "ingest",
                "module": "file.read",
                "params": {"path": json_path},
            },
            {
                "id": "parse",
                "module": "data.json.parse",
                "params": {"json_string": SAMPLE_JSON},
            },
            {
                "id": "report",
                "module": "data.text.template",
                "params": {
                    "template": "Team Report: {team}\n===========\nMembers: {count}\nRoles: {roles}\nGenerated by: flyto-core",
                    "variables": {
                        "team": "Alpha",
                        "count": "3",
                        "roles": "Engineer, Designer, Manager",
                    },
                },
            },
            {
                "id": "status",
                "module": "string.uppercase",
                "params": {"text": "pipeline complete — all steps passed"},
            },
            {
                "id": "export",
                "module": "file.write",
                "params": {
                    "path": report_path,
                    "content": "Team Report: Alpha\n===========\n3 members\nStatus: COMPLETE",
                },
            },
        ],
    }

    execution_id = f"exec_{int(time.time())}"
    hooks = APIEvidenceHooks(evidence_store, execution_id)

    engine = WorkflowEngine(
        workflow=workflow,
        params={},
        hooks=hooks,
        enable_trace=True,
    )

    t0 = time.time()
    await engine.execute()
    elapsed = int((time.time() - t0) * 1000)

    trace = engine.get_execution_trace_dict()
    for step in trace.get("steps", []):
        items = step.get("output", {}).get("items", [[]])
        first = items[0][0] if items and items[0] else {}
        display = summarize_output(first)
        ok(f"{step['stepId']:10s} {D}({step['moduleId']}){X} → {B}{display}{X}")

    print(f"\n  {G}Completed{X} in {Y}{elapsed}ms{X} · {C}{execution_id}{X}")

    # Save workflow for replay
    exec_dir = evidence_store.get_execution_dir(execution_id)
    with open(exec_dir / "workflow.json", "w") as f:
        json.dump(workflow, f, indent=2)

    # ── 2. Show execution trace ──────────────────────────────────

    header("2. Execution Trace")

    fmt = "  {:<4s} {:<10s} {:<22s} {:<9s} {:<s}"
    print(fmt.format("#", "Step", "Module", "Status", "Output"))
    print(f"  {'─' * 60}")
    for i, step in enumerate(trace.get("steps", [])):
        items = step.get("output", {}).get("items", [[]])
        first = items[0][0] if items and items[0] else {}
        display = summarize_output(first)[:30]
        color = G if step["status"] == "success" else R
        print(fmt.format(
            str(i + 1),
            step["stepId"],
            step["moduleId"],
            f"{color}{step['status']}{X}",
            display,
        ))

    # ── 3. Show evidence (state snapshots) ───────────────────────

    header("3. Evidence — Context Snapshots")

    await asyncio.sleep(0.2)  # let async evidence writes flush
    evidence_list = await evidence_store.load_evidence(execution_id)

    for ev in evidence_list:
        before = list(ev.context_before.keys()) if ev.context_before else []
        after = list(ev.context_after.keys()) if ev.context_after else []
        new_keys = [k for k in after if k not in before]

        print(f"\n  {B}{ev.step_id}{X} {D}· {ev.module_id}{X}")
        print(f"    context_before: {D}{{ {', '.join(before) if before else ''} }}{X}")
        print(f"    context_after:  {{ {', '.join(after)} }}")
        if new_keys:
            print(f"    {G}+ {', '.join(new_keys)}{X}")

    info(f"\n  Evidence on disk: {evidence_path / execution_id}/")

    # ── 4. Replay from step 3 ────────────────────────────────────

    header("4. Replay From Step 3")

    from core.engine.replay import create_replay_manager

    replay_manager = create_replay_manager(evidence_path)

    async def workflow_executor(workflow, context, start_step, end_step=None, skip_steps=None, breakpoints=None):
        steps = workflow.get("steps", [])
        start_idx = None
        for i, s in enumerate(steps):
            if s.get("id") == start_step:
                start_idx = i
                break

        replay_id = f"replay_{int(time.time())}"
        replay_hooks = APIEvidenceHooks(evidence_store, replay_id)
        replay_engine = WorkflowEngine(
            workflow=workflow,
            params={},
            hooks=replay_hooks,
            enable_trace=True,
            start_step=start_idx,
            initial_context=context,
        )
        await replay_engine.execute()
        return {
            "ok": True,
            "execution_id": replay_id,
            "steps_executed": len(replay_engine.execution_log),
        }

    result = await replay_manager.replay_from_step(
        execution_id=execution_id,
        step_id="report",
        workflow_executor=workflow_executor,
    )

    if result.ok:
        ok(f"Replay from {B}report{X} (step 3)")
        ok(f"Re-ran {Y}{result.steps_executed}{X} steps (skipped ingest + parse)")
        ok(f"Replay ID: {C}{result.execution_id}{X}")
        info(f"Steps 1-2 skipped. Context loaded from evidence.")
    else:
        fail(f"Replay failed: {result.error}")

    # ── Done ─────────────────────────────────────────────────────

    header("Done")

    print(f"""
  {B}What just happened:{X}

  {G}✓{X} Data pipeline: file → JSON parse → template → format → export
  {G}✓{X} Execution trace captured for every step
  {G}✓{X} Evidence (context_before/after) stored per step
  {G}✓{X} Replayed from step 3 without re-running steps 1-2

  {B}Next steps:{X}

    {D}# Use with your AI (MCP):{X}
    pip install flyto-core
    {D}# Add to Claude Code / Cursor / Windsurf config (see README){X}

    {D}# Or start the HTTP API:{X}
    flyto serve
    curl localhost:8333/v1/modules  {D}# discover all 329 modules{X}

  {P}github.com/flytohub/flyto-core{X}
""")


if __name__ == "__main__":
    asyncio.run(main())
