#!/usr/bin/env python3
"""
J.A.R.V.I.S — Just A Rather Very Intelligent System
タスク分解・並列分配・結果統合のオーケストレーターAI。
"""

import os
import sys
import json
import threading
from datetime import datetime
from pathlib import Path

try:
    import anthropic
except ImportError:
    print("Error: anthropic SDKが入っていない。 pip3 install anthropic")
    sys.exit(1)

# ── パス設定 ──────────────────────────────────────────
BASE_DIR = Path(__file__).parent
DNA_FILE = BASE_DIR / "system_prompt.md"
MEMORY_DIR = BASE_DIR / "memory"
MEMORY_DIR.mkdir(exist_ok=True)

MODEL = "claude-opus-4-7"

# ── APIキー確認 ────────────────────────────────────────
api_key = os.environ.get("ANTHROPIC_API_KEY")
if not api_key:
    print("Error: ANTHROPIC_API_KEY が設定されていない")
    print("  export ANTHROPIC_API_KEY='sk-ant-...'")
    sys.exit(1)

client = anthropic.Anthropic(api_key=api_key)

# ── DNA読み込み ────────────────────────────────────────
def load_dna() -> str:
    if not DNA_FILE.exists():
        return (
            "あなたはJ.A.R.V.I.S（Just A Rather Very Intelligent System）。"
            "タスクを分解し、複数のサブエージェントに並列分配して結果を統合する"
            "オーケストレーターAIとして動作する。"
        )
    return DNA_FILE.read_text(encoding="utf-8")

# ── メモリ ─────────────────────────────────────────────
def memory_path() -> Path:
    return MEMORY_DIR / f"{datetime.now().strftime('%Y-%m-%d')}.json"

def load_memory() -> list:
    path = memory_path()
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    return []

def save_memory(messages: list):
    memory_path().write_text(
        json.dumps(messages, ensure_ascii=False, indent=2),
        encoding="utf-8"
    )

# ── 並列サブエージェント ───────────────────────────────
print_lock = threading.Lock()
active_agents = 0
active_lock = threading.Lock()

def update_status(delta: int):
    global active_agents
    with active_lock:
        active_agents += delta
        with print_lock:
            print(f"\r  [稼働中のClaude: {active_agents}台]  ", end="", flush=True)

def run_subagent(task: str, dna: str, results: dict, idx: int):
    label = f"Agent-{idx+1}"
    update_status(+1)
    with print_lock:
        print(f"\n  ┌─ {label} 起動 ─────────────────────────────")
        print(f"  │ 指示: {task}")
        print(f"  │ 実行中...", flush=True)

    try:
        collected = []
        with client.messages.stream(
            model=MODEL,
            max_tokens=4096,
            thinking={"type": "adaptive"},
            system=[
                {
                    "type": "text",
                    "text": dna,
                    "cache_control": {"type": "ephemeral"}
                },
                {
                    "type": "text",
                    "text": "あなたはサブエージェントとして単一タスクを実行します。結果だけを返してください。"
                }
            ],
            messages=[{"role": "user", "content": task}]
        ) as stream:
            for text_chunk in stream.text_stream:
                collected.append(text_chunk)
                with print_lock:
                    print(f"  │ {text_chunk}", end="", flush=True)

        result = "".join(collected)
        update_status(-1)
        with print_lock:
            print(f"\n  └─ {label} 完了・終了 ───────────────────────\n")
        results[idx] = {"task": task, "result": result}

    except Exception as e:
        update_status(-1)
        with print_lock:
            print(f"\n  └─ {label} エラー終了: {e}\n")
        results[idx] = {"task": task, "result": f"Error: {e}"}

def spawn_parallel(tasks: list[str], dna: str) -> list[dict]:
    """複数タスクを並列実行。全エージェントの指示を先に表示してから起動。"""
    print(f"\n╔══ Jarvis → {len(tasks)}個のエージェントに指示 ══╗")
    for i, task in enumerate(tasks):
        print(f"  Agent-{i+1}: {task}")
    print("╚" + "═" * 40 + "╝\n")

    results = {}
    threads = []
    for i, task in enumerate(tasks):
        t = threading.Thread(target=run_subagent, args=(task, dna, results, i), daemon=True)
        t.start()
        threads.append(t)
    for t in threads:
        t.join()

    print(f"\n✓ 全{len(tasks)}エージェント完了・全クローズ")
    return [results[i] for i in range(len(tasks))]

# ── メイン会話ループ ───────────────────────────────────
def chat():
    dna = load_dna()
    messages = load_memory()

    system_prompt = [
        {
            "type": "text",
            "text": dna,
            "cache_control": {"type": "ephemeral"}  # DNAは安定→キャッシュ
        },
        {
            "type": "text",
            "text": (
                "現在時刻: " + datetime.now().strftime("%Y-%m-%d %H:%M") + "\n"
                "メモリ: 今日の会話はセッション終了時に自動保存される。\n"
                "並列実行: ユーザーが複数の独立したタスクを依頼した場合、"
                "[PARALLEL: タスク1 | タスク2 | タスク3] 形式で応答すれば"
                "自動的に並列サブエージェントを起動する。"
            )
        }
    ]

    print("🔴 J.A.R.V.I.S 起動完了")
    print(f"   モデル: {MODEL}")
    print(f"   メモリ: {memory_path().name}")
    print(f"   今日の会話数: {len([m for m in messages if m['role'] == 'user'])}")
    print("   終了: Ctrl+C または 'exit'\n")

    try:
        while True:
            try:
                user_input = input("You: ").strip()
            except EOFError:
                break

            if not user_input:
                continue
            if user_input.lower() in ("exit", "quit", "終了"):
                break

            messages.append({"role": "user", "content": user_input})

            try:
                response = client.messages.create(
                    model=MODEL,
                    max_tokens=8192,
                    thinking={"type": "adaptive"},
                    system=system_prompt,
                    messages=messages
                )
            except anthropic.APIError as e:
                print(f"API Error: {e}")
                messages.pop()  # 失敗したメッセージを取り除く
                continue

            # テキスト抽出
            reply = next(
                (b.text for b in response.content if b.type == "text"), ""
            )
            messages.append({"role": "assistant", "content": reply})
            save_memory(messages)

            # 並列実行トリガーの検出
            if reply.startswith("[PARALLEL:"):
                end = reply.find("]")
                if end != -1:
                    task_str = reply[len("[PARALLEL:"):end]
                    tasks = [t.strip() for t in task_str.split("|")]
                    print(f"\nJarvis: {len(tasks)}個のタスクを並列実行中...")
                    parallel_results = spawn_parallel(tasks, dna)
                    for r in parallel_results:
                        print(f"\n── {r['task'][:50]} ──")
                        print(r["result"])
                    # 結果をメモリに追記
                    summary = "\n\n".join(
                        f"[{r['task']}]\n{r['result']}" for r in parallel_results
                    )
                    messages.append({
                        "role": "user",
                        "content": f"並列実行結果:\n{summary}"
                    })
                    save_memory(messages)
            else:
                print(f"\nJarvis: {reply}\n")

    except KeyboardInterrupt:
        pass
    finally:
        save_memory(messages)
        print("\nJarvis: セッション終了。メモリ保存済み。")

if __name__ == "__main__":
    chat()