#!/usr/bin/env python3
import requests, subprocess, json, os

API_BASE = "http://localhost:8000/v1"
MODEL = "qwen3.6-35b"
# WORKDIR = os.path.expanduser("~")  # /home/windpacer

def run_shell_command(cmd: str, cwd: str = None) -> dict:
    """쉘 명령어 실행 + 메타데이터 반환"""
    effective_cwd = cwd if cwd else os.getcwd()
    try:
        result = subprocess.run(
            cmd, shell=True, capture_output=True, text=True, 
            timeout=30, cwd=effective_cwd, executable='/bin/bash'
        )
        return {
            "success": result.returncode == 0,
            "stdout": result.stdout.strip(),
            "stderr": result.stderr.strip(),
            "cwd": effective_cwd
        }
    except Exception as e:
        return {"success": False, "error": str(e)}

def chat_with_tool(user_msg: str, command: str):
    tools = [{
        "type": "function",
        "function": {
            "name": "runShellCommand",
            "description": "Execute shell command on DGX Spark (ARM64 Ubuntu)",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {"type": "string", "description": "Command to execute"}
                },
                "required": ["command"]
            }
        }
    }]
    
    # 1 차: 모델이 도구 호출
    print(f"🤖 질문: {user_msg}")
    resp1 = requests.post(f"{API_BASE}/chat/completions", json={
        "model": MODEL,
        "messages": [{"role": "user", "content": user_msg}],
        "tools": tools,
        "tool_choice": "auto",
        "stream": False
    }, timeout=120).json()
    
    msg1 = resp1["choices"][0]["message"]
    
    # tool_calls 확인
    if not msg1.get("tool_calls"):
        print(f"⚠️ 도구 호출 없음: {msg1.get('content', 'No content')[:200]}...")
        return
    
    # 도구 실행
    for tool in msg1["tool_calls"]:
        if tool["function"]["name"] == "runShellCommand":
            args = json.loads(tool["function"]["arguments"])
            actual_cmd = args.get("command", command)
            print(f"🔧 실행: {actual_cmd}")
            
            output = run_shell_command(actual_cmd)
            print(f"📤 결과: {json.dumps(output, ensure_ascii=False, indent=2)}")
            
            # 2 차: 결과 피드백
            messages = [
                {"role": "user", "content": user_msg},
                msg1,
                {
                    "role": "tool",
                    "content": json.dumps(output, ensure_ascii=False),
                    "tool_call_id": tool["id"]
                }
            ]
            
            resp2 = requests.post(f"{API_BASE}/chat/completions", json={
                "model": MODEL,
                "messages": messages,
                "tools": tools,
                "tool_choice": "none",
                "stream": False
            }, timeout=120).json()
            
            final = resp2["choices"][0]["message"]["content"]
            print(f"\n✅ 최종 응답:\n{final}")
            return final

if __name__ == "__main__":
    # 테스트 1: pwd
    chat_with_tool("현재 작업 디렉토리가 어디야?", "pwd")
    print("\n" + "="*60 + "\n")
    # 테스트 2: ls + 파일 확인
    chat_with_tool("add.py 파일이 현재 디렉토리에 있어?", "ls -l add.py")