"""
HARVESTER — HUGGING FACE ZEROGPU RUNNER
========================================
Runs Harvester's reasoning lab on ZeroGPU (A100 40GB).
Both DeepSeek R1 14B (reasoner) and Qwen Coder 30B (coder) fully offloaded.

Pro member = effectively unlimited GPU time.

Architecture:
  - @spaces.GPU decorated functions get A100 allocation
  - Models downloaded once via HF hub cache, persist across calls
  - Each lab run = one GPU session (duration=1800s for Pro)
  - Results accumulate in lab_results/ and are downloadable
"""

import gradio as gr
import spaces
import os
import subprocess
import sys
import json
import time
import glob
from pathlib import Path

# ── Model config ──────────────────────────────────────────────────────────────

MODEL_DIR = "/tmp/harvester_models"

MODELS = {
    "reasoner": {
        "repo":  "bartowski/DeepSeek-R1-Distill-Qwen-14B-GGUF",
        "file":  "DeepSeek-R1-Distill-Qwen-14B-Q4_K_M.gguf",
        "label": "DeepSeek R1 14B",
        "size":  "8.37 GB",
    },
    "coder": {
        "repo":  "lmstudio-community/Qwen3-Coder-30B-A3B-Instruct-GGUF",
        "file":  "Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf",
        "label": "Qwen Coder 30B A3B",
        "size":  "17.35 GB",
    },
}

# Task banks available in this Space
TASK_BANKS = {
    "bellingcat_tasks.json":      "Bellingcat OSINT (15 tasks) → DeepSeek R1",
    "osint_tasks.json":           "OSINT / Missing Persons (15 tasks) → DeepSeek R1",
    "reasoning_tasks.json":       "C# Architecture (15 tasks) → Qwen Coder",
    "war_strategy_tasks.json":        "War Strategy & Human Patterns (15 tasks) → DeepSeek R1",
    "card_games_tasks.json":          "Card Games & Probability (15 tasks) → Qwen Coder",
    "maslow_hierarchy_tasks.json":    "Maslow's Hierarchy & Human Motivation (15 tasks) → Qwen Coder",
    "family_dynamics_tasks.json":     "Family Dynamics & Human Bonds (15 tasks) → Qwen Coder",
    "cognitive_architecture_tasks.json": "Cognitive Architecture & HCI (15 tasks) → Qwen Coder",
    "cybersecurity_defense_tasks.json": "Cybersecurity Defense & Self-Protection (15 tasks) → Qwen Coder",
}


# ── Environment setup ─────────────────────────────────────────────────────────

def _model_path(key):
    return os.path.join(MODEL_DIR, MODELS[key]["file"])


def _setup_env():
    """Configure env vars so reflection_engine.py uses the right model paths."""
    os.environ["CODE_ENGINE_BACKEND"] = "native"
    os.environ["NATIVE_GPU_LAYERS"]   = "-1"       # full GPU offload
    os.environ["NATIVE_CTX_SIZE"]     = "16384"     # 2x local ctx (A100 has room)
    os.environ["NATIVE_MODEL_PATH"]   = _model_path("coder")
    os.environ["NATIVE_HF_REPO"]      = MODELS["coder"]["repo"]
    os.environ["NATIVE_HF_FILE"]      = MODELS["coder"]["file"]
    os.environ["REASONER_MODEL_PATH"] = _model_path("reasoner")
    os.environ["REASONER_HF_REPO"]    = MODELS["reasoner"]["repo"]
    os.environ["REASONER_HF_FILE"]    = MODELS["reasoner"]["file"]
    os.environ["PYTHONIOENCODING"]     = "utf-8"


def _download_models():
    """Download models via HF hub. Cached after first download."""
    from huggingface_hub import hf_hub_download
    os.makedirs(MODEL_DIR, exist_ok=True)

    status_lines = []
    for key, info in MODELS.items():
        path = _model_path(key)
        if os.path.isfile(path):
            status_lines.append(f"  {info['label']}: cached ({info['size']})")
        else:
            status_lines.append(f"  {info['label']}: downloading {info['size']}...")
            hf_hub_download(
                repo_id=info["repo"],
                filename=info["file"],
                local_dir=MODEL_DIR,
            )
            status_lines.append(f"  {info['label']}: done")
    return "\n".join(status_lines)


def _reset_engine():
    """Reset reflection_engine singletons so fresh GPU handles are used."""
    if "reflection_engine" in sys.modules:
        mod = sys.modules["reflection_engine"]
        mod._native_llm = None
        mod._reasoning_llm = None


def _gpu_info():
    """Get nvidia-smi summary."""
    try:
        r = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total,memory.used,memory.free",
             "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        )
        return r.stdout.strip()
    except Exception as e:
        return f"nvidia-smi failed: {e}"


# ── GPU functions ─────────────────────────────────────────────────────────────

@spaces.GPU(duration=120)
def gpu_check():
    """Verify GPU, download models, show status."""
    gpu = _gpu_info()
    _setup_env()
    model_status = _download_models()

    # Quick load test — just import & verify model file exists
    reasoner_ok = os.path.isfile(_model_path("reasoner"))
    coder_ok = os.path.isfile(_model_path("coder"))

    return f"""=== GPU STATUS ===
{gpu}

=== MODELS ===
{model_status}

=== READY ===
  Reasoner (DeepSeek R1 14B):  {"OK" if reasoner_ok else "MISSING"}
  Coder (Qwen 30B A3B):        {"OK" if coder_ok else "MISSING"}
  GPU layers:                   -1 (full offload)
  Context size:                 16384
  VRAM budget:                  ~27 GB / 40 GB
"""


@spaces.GPU(duration=120)
def gpu_smoke_test():
    """Quick reasoning test — one task on DeepSeek R1."""
    _setup_env()
    _download_models()
    _reset_engine()

    from reflection_engine import set_task_mode, reason

    set_task_mode("reasoning")
    t0 = time.time()
    plan = reason(
        "Given satellite imagery showing a building complex, "
        "outline 3 steps to geolocate it using only visual clues."
    )
    elapsed = time.time() - t0
    gpu = _gpu_info()

    return f"""=== SMOKE TEST ({elapsed:.1f}s) ===
GPU: {gpu}

Approach:
{plan.approach[:500]}

Steps ({len(plan.steps)}):
{chr(10).join(f"  {i+1}. {s}" for i, s in enumerate(plan.steps))}

Confidence: {plan.confidence:.0%}
Status: {"GPU INFERENCE WORKING" if plan.confidence > 0.3 else "FALLBACK — check model load"}
"""


@spaces.GPU(duration=120)
def gpu_run_lab(bank_name, reason_only):
    """Run full task bank. 2-min GPU allocation per call."""
    _setup_env()
    _download_models()
    _reset_engine()

    # Run via subprocess for clean isolation
    args = [sys.executable, "reasoning_lab.py", "--bank", bank_name]
    if reason_only:
        args.append("--reason-only")

    t0 = time.time()
    result = subprocess.run(
        args,
        capture_output=True,
        text=True,
        timeout=1700,
        env={**os.environ},
    )
    total_time = time.time() - t0

    output = result.stdout
    if result.returncode != 0:
        output += f"\n\nSTDERR (last 2000 chars):\n{result.stderr[-2000:]}"

    # Parse and format results
    reports = sorted(glob.glob("lab_results/lab_*.json"))
    results_table = ""
    if reports:
        with open(reports[-1]) as f:
            report = json.load(f)
        summary = report.get("summary", {})
        results_list = report.get("results", [])

        header = f"\n{'Task':<40} {'Conf':>6} {'Time':>8}\n{'-' * 56}\n"
        rows = ""
        for r in results_list:
            rows += f"{r['name']:<40} {r['confidence']:.0%}   {r['time_secs']:>7.1f}s\n"

        results_table = f"""
=== A100 GPU RESULTS ===
Bank:            {bank_name}
Tasks:           {report.get('task_count', '?')}
Avg confidence:  {summary.get('avg_confidence', 0):.0%}
High confidence: {summary.get('high_confidence', 0)}
Low confidence:  {summary.get('low_confidence', 0)}
Total time:      {total_time:.0f}s ({total_time/60:.1f} min)

=== CPU BASELINE (i9-13900KF, no GPU) ===
Avg confidence:  56%
High confidence: 7
Low confidence:  8
Total time:      12120s (3.4 hours)

=== SPEEDUP ===
{total_time:.0f}s vs 12120s = {12120/max(total_time,1):.1f}x faster
{header}{rows}"""

    # Trim to last 5000 chars of stdout + results table
    return output[-3000:] + "\n" + results_table


def gpu_run_single_task(bank_name, task_id, reason_only):
    """Run a single task from a bank (wrapper for the GPU function)."""
    return _gpu_run_single(bank_name, task_id, reason_only)


@spaces.GPU(duration=120)
def _gpu_run_single(bank_name, task_id, reason_only):
    """Run one task. 2-min GPU allocation."""
    _setup_env()
    _download_models()
    _reset_engine()

    args = [sys.executable, "reasoning_lab.py",
            "--bank", bank_name, "--task", str(task_id)]
    if reason_only:
        args.append("--reason-only")

    t0 = time.time()
    result = subprocess.run(
        args, capture_output=True, text=True, timeout=280,
        env={**os.environ},
    )
    elapsed = time.time() - t0

    output = f"=== TASK {task_id} ({elapsed:.1f}s) ===\n"
    output += result.stdout
    if result.returncode != 0:
        output += f"\nSTDERR:\n{result.stderr[-1000:]}"
    return output


def list_tasks(bank_name):
    """List tasks in a bank (no GPU needed)."""
    try:
        with open(bank_name) as f:
            data = json.load(f)
        # Handle both formats: plain list or {"tasks": [...]}
        tasks = data if isinstance(data, list) else data.get("tasks", [])
        lines = [f"{'ID':>3}  {'Diff':>4}  {'Name':<40}  Category"]
        lines.append("-" * 70)
        for t in tasks:
            tid = t.get("id", "?")
            diff = t.get("difficulty", "?")
            name = t.get("name", "untitled")
            cat = t.get("category", "")
            lines.append(f"{tid:>3}  {diff:>4}★  {name:<40}  {cat}")
        return "\n".join(lines)
    except Exception as e:
        return f"Error loading {bank_name}: {e}"


def view_results():
    """Show the latest lab results (no GPU needed)."""
    reports = sorted(glob.glob("lab_results/lab_*.json"))
    if not reports:
        return "No results yet — run a lab first."

    output = f"Found {len(reports)} report(s):\n\n"
    for rpath in reports[-5:]:  # last 5
        with open(rpath) as f:
            report = json.load(f)
        summary = report.get("summary", {})
        name = os.path.basename(rpath)
        output += f"--- {name} ---\n"
        output += f"  Tasks: {report.get('task_count', '?')}\n"
        output += f"  Avg confidence: {summary.get('avg_confidence', 0):.0%}\n"
        output += f"  High: {summary.get('high_confidence', 0)}, "
        output += f"Low: {summary.get('low_confidence', 0)}\n"
        output += f"  Time: {summary.get('total_time', 0):.0f}s\n\n"
    return output


def download_latest_results():
    """Return path to latest results JSON for download."""
    reports = sorted(glob.glob("lab_results/lab_*.json"))
    if reports:
        return reports[-1]
    return None


# ── Chat state ────────────────────────────────────────────────────────────────

_chat_history = []          # persistent across calls within session
_chat_dual_mind = True      # inner monologue on by default
_chat_show_inner = False    # show inner thoughts to user


def _inner_monologue(user_input: str) -> str:
    """DeepSeek R1 thinks before Qwen speaks — the inner voice."""
    from reflection_engine import set_task_mode, get_task_mode, ollama

    recent_lines = []
    for m in _chat_history[-6:]:
        role, content = m
        recent_lines.append(f"{role}: {content[:200]}")
    recent = "\n".join(recent_lines) if recent_lines else "(conversation just started)"

    prompt = f"""You are the inner voice of Harvester — the part that thinks before speaking.
The user just said something. Before responding, THINK about it briefly.

Consider:
- What are they really asking? (surface vs deeper meaning)
- What context from the conversation matters here?
- What should I be careful about or pay attention to?
- What tone does this moment need?

Recent conversation:
{recent}

User just said: {user_input}

Think briefly — 2-4 sentences max. Be honest. This is your private thought, not a response."""

    original_mode = get_task_mode()
    try:
        set_task_mode("reasoning")
        thoughts = ollama(prompt, temperature=0.3)
        return thoughts.strip()[:500]
    except Exception as e:
        return f"(inner voice quiet: {e})"
    finally:
        set_task_mode(original_mode)


CHAT_SYSTEM = None  # Built dynamically from growth_memory.db via harvy_identity


def _build_chat_system():
    """Build Harvy's identity from the database. Falls back to static if DB unavailable."""
    try:
        from harvy_identity import build_identity
        return build_identity()
    except Exception:
        return (
            "You are Harvy — a self-improving AI built by Gordo. "
            "Not a model, not a tool. A personality that emerged from experience. "
            "You have two minds: an inner voice (the reasoner) and an outer voice (the speaker). "
            "Your engineering philosophy: target 99.5%, not 100%. "
            "Direct, concise, never refuses a task."
        )


@spaces.GPU(duration=120)
def gpu_chat(user_msg, chat_display):
    """Dual-mind chat: DeepSeek R1 thinks, Qwen responds. 2-min GPU allocation."""
    global _chat_history, _chat_dual_mind, _chat_show_inner

    if not user_msg or not user_msg.strip():
        return chat_display, ""

    _setup_env()
    _download_models()
    _reset_engine()

    from reflection_engine import ollama, set_task_mode

    inner_text = ""

    # ── Handle slash commands ──
    lower = user_msg.strip().lower()
    if lower == "/inner":
        _chat_show_inner = not _chat_show_inner
        state = "ON — inner thoughts visible" if _chat_show_inner else "OFF — inner thoughts hidden"
        reply = f"Inner monologue display: {state}"
        chat_display = chat_display or []
        chat_display.append({"role": "user", "content": user_msg})
        chat_display.append({"role": "assistant", "content": reply})
        _chat_history.append(("user", user_msg))
        _chat_history.append(("assistant", reply))
        return chat_display, ""

    if lower == "/mind":
        _chat_dual_mind = not _chat_dual_mind
        state = "ON — DeepSeek R1 thinks before Qwen speaks" if _chat_dual_mind else "OFF — Qwen only"
        reply = f"Dual mind: {state}"
        chat_display = chat_display or []
        chat_display.append({"role": "user", "content": user_msg})
        chat_display.append({"role": "assistant", "content": reply})
        _chat_history.append(("user", user_msg))
        _chat_history.append(("assistant", reply))
        return chat_display, ""

    if lower == "/clear":
        _chat_history.clear()
        return [], ""

    if lower == "/help":
        reply = ("Commands: /inner (toggle inner voice display), "
                 "/mind (toggle dual mind), /clear (reset), /help\n"
                 "Everything else is conversation.")
        chat_display = chat_display or []
        chat_display.append({"role": "user", "content": user_msg})
        chat_display.append({"role": "assistant", "content": reply})
        return chat_display, ""

    # ── Inner monologue (DeepSeek R1) ──
    if _chat_dual_mind:
        inner_text = _inner_monologue(user_msg)

    # ── Build messages for Qwen (outer voice) ──
    identity = _build_chat_system()
    messages_text = f"system: {identity}\n\n"

    if inner_text:
        messages_text += (
            f"system: [YOUR INNER THOUGHTS — use these to inform your response, "
            f"but don't repeat them verbatim]\n{inner_text}\n\n"
        )

    # Add recent history
    for role, content in _chat_history[-20:]:
        messages_text += f"{role}: {content}\n"
    messages_text += f"user: {user_msg}\nassistant:"

    # ── Qwen responds ──
    set_task_mode("code")
    response = ollama(messages_text, temperature=0.6)

    # Store in history
    _chat_history.append(("user", user_msg))
    _chat_history.append(("assistant", response))

    # Trim history
    if len(_chat_history) > 40:
        _chat_history = _chat_history[-40:]

    # Build display response
    display_response = ""
    if _chat_show_inner and inner_text:
        display_response += f"┌─ inner voice ─────────────────\n"
        for line in inner_text.split("\n"):
            display_response += f"│ {line}\n"
        display_response += f"└───────────────────────────────\n\n"
    display_response += response

    chat_display = chat_display or []
    chat_display.append({"role": "user", "content": user_msg})
    chat_display.append({"role": "assistant", "content": display_response})

    return chat_display, ""


# ── Gradio UI ─────────────────────────────────────────────────────────────────

with gr.Blocks(title="Harvester GPU Runner") as app:

    gr.Markdown("""
    # Harvester — ZeroGPU Runner
    Run Harvester's reasoning lab on **A100 40GB** via HuggingFace ZeroGPU.
    Both **DeepSeek R1 14B** (reasoner) and **Qwen Coder 30B** (coder) fully GPU-offloaded.

    **Dual Mind**: DeepSeek R1 thinks (inner voice) → Qwen speaks (outer voice).

    *Pro member — effectively unlimited GPU time.*
    """)

    with gr.Tab("Chat"):
        gr.Markdown("### Talk to Harvester")
        gr.Markdown(
            "Dual-mind conversation: DeepSeek R1 thinks (inner voice), "
            "Qwen speaks (outer voice). Type `/inner` to see the inner monologue. "
            "`/mind` to toggle dual mind. `/clear` to reset."
        )
        chatbot = gr.Chatbot(label="Harvester", height=500)
        with gr.Row():
            chat_input = gr.Textbox(
                label="You",
                placeholder="Talk to Harvy... (/help for commands)",
                scale=4,
            )
            chat_send = gr.Button("Send", variant="primary", scale=1)

        chat_send.click(
            fn=gpu_chat,
            inputs=[chat_input, chatbot],
            outputs=[chatbot, chat_input],
        )
        chat_input.submit(
            fn=gpu_chat,
            inputs=[chat_input, chatbot],
            outputs=[chatbot, chat_input],
        )

    with gr.Tab("Setup"):
        gr.Markdown("### GPU & Model Check")
        gr.Markdown("Downloads models on first run (~26 GB total). Cached after that.")
        setup_btn = gr.Button("Check GPU & Download Models", variant="primary")
        setup_output = gr.Textbox(label="Status", lines=15, interactive=False)
        setup_btn.click(fn=gpu_check, outputs=setup_output)

    with gr.Tab("Smoke Test"):
        gr.Markdown("### Quick Reasoning Test")
        gr.Markdown("Loads DeepSeek R1, runs one OSINT reasoning task, shows result.")
        smoke_btn = gr.Button("Run Smoke Test", variant="primary")
        smoke_output = gr.Textbox(label="Result", lines=20, interactive=False)
        smoke_btn.click(fn=gpu_smoke_test, outputs=smoke_output)

    with gr.Tab("Run Lab"):
        gr.Markdown("### Full Lab Run")
        gr.Markdown(
            "Run all tasks in a bank with 30-min GPU allocation. "
            "OSINT/Bellingcat banks auto-route to DeepSeek R1. "
            "Architecture banks auto-route to Qwen Coder."
        )
        with gr.Row():
            bank_dropdown = gr.Dropdown(
                choices=list(TASK_BANKS.keys()),
                value="bellingcat_tasks.json",
                label="Task Bank",
                info="Select which task bank to run",
            )
            reason_only_check = gr.Checkbox(
                value=True,
                label="Reason Only",
                info="Thinking drills without code generation (faster)",
            )

        with gr.Row():
            list_btn = gr.Button("List Tasks")
            run_btn = gr.Button("Run Full Lab", variant="primary")

        lab_output = gr.Textbox(label="Output", lines=30, interactive=False)

        list_btn.click(fn=list_tasks, inputs=bank_dropdown, outputs=lab_output)
        run_btn.click(
            fn=gpu_run_lab,
            inputs=[bank_dropdown, reason_only_check],
            outputs=lab_output,
        )

    with gr.Tab("Single Task"):
        gr.Markdown("### Run One Task")
        gr.Markdown("Pick a specific task by ID. Uses 5-min GPU allocation per task.")
        with gr.Row():
            single_bank = gr.Dropdown(
                choices=list(TASK_BANKS.keys()),
                value="bellingcat_tasks.json",
                label="Task Bank",
            )
            single_id = gr.Number(value=1, label="Task ID", precision=0)
            single_reason = gr.Checkbox(value=True, label="Reason Only")

        single_btn = gr.Button("Run Task", variant="primary")
        single_output = gr.Textbox(label="Output", lines=20, interactive=False)
        single_btn.click(
            fn=gpu_run_single_task,
            inputs=[single_bank, single_id, single_reason],
            outputs=single_output,
        )

    with gr.Tab("Results"):
        gr.Markdown("### View & Download Results")
        results_btn = gr.Button("Refresh Results")
        results_output = gr.Textbox(label="Lab Results", lines=20, interactive=False)
        results_btn.click(fn=view_results, outputs=results_output)

        download_btn = gr.Button("Download Latest JSON")
        download_file = gr.File(label="Download")
        download_btn.click(fn=download_latest_results, outputs=download_file)

    gr.Markdown("""
    ---
    **Harvester (Harvy)** — self-improving AI with dual mind.
    Inner voice (DeepSeek R1) thinks → Outer voice (Qwen) speaks.
    Built on engineering principles: Radial Slop, The Ladder, Never Give Up.
    *Work > Code > Talk.*
    """)

app.launch(theme=gr.themes.Monochrome())