import asyncio, json, os, subprocess, sys, tempfile, re, textwrap
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse, quote

import uvicorn
import aiohttp
import aiofiles
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI(title="OpenDAN Web")

LLAMA_CHAT = "http://127.0.0.1:8081"
LLAMA_EMBED = "http://127.0.0.1:8082"
MYAI = Path.home() / "myai"
MAX_HISTORY = 50
conversations: dict[str, list[dict]] = {}
proposals: dict[str, dict] = {}

TOOL_DEFINITIONS = [
    {
        "type": "function",
        "function": {
            "name": "system.shell.exec",
            "description": "Execute a shell command in Linux bash. Use for running programs, system administration, file operations.",
            "parameters": {
                "type": "object",
                "properties": {"command": {"type": "string", "description": "The bash command to execute"}},
                "required": ["command"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "system.code_interpreter",
            "description": "Execute Python code in an isolated environment. Use for data analysis, calculations, scripting.",
            "parameters": {
                "type": "object",
                "properties": {"code": {"type": "string", "description": "Python code to execute"}},
                "required": ["code"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "agent.workspace.read_file",
            "description": "Read a file from the user's filesystem.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string", "description": "Absolute path to the file"}},
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "agent.workspace.write_file",
            "description": "Write content to a file. Creates parent directories if needed.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "Absolute path to the file"},
                    "content": {"type": "string", "description": "Content to write"}
                },
                "required": ["path", "content"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "agent.workspace.list_dir",
            "description": "List files and directories in a given path.",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string", "description": "Directory path to list"}},
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "web.search",
            "description": "Search the web using DuckDuckGo. Use for finding current information.",
            "parameters": {
                "type": "object",
                "properties": {"query": {"type": "string", "description": "Search query"}},
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "system.now",
            "description": "Get the current date and time.",
            "parameters": {"type": "object", "properties": {}}
        }
    },
]

TOOL_IMPLS = {}

async def tool_shell_exec(command: str) -> dict:
    try:
        proc = await asyncio.create_subprocess_shell(
            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            cwd=str(MYAI))
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=30)
        out = stdout.decode(errors="replace").strip()
        err = stderr.decode(errors="replace").strip()
        return {"status": "success" if proc.returncode == 0 else "error",
                "stdout": out[:10000], "stderr": err[:5000],
                "returncode": proc.returncode}
    except asyncio.TimeoutError:
        return {"status": "error", "error": "Command timed out after 30s"}
    except Exception as e:
        return {"status": "error", "error": str(e)}

async def tool_code_interpreter(code: str) -> dict:
    tmp_dir = MYAI / "tmp_code"
    tmp_dir.mkdir(parents=True, exist_ok=True)
    script = tmp_dir / "_web_exec.py"
    async with aiofiles.open(script, "w") as f:
        await f.write(code)
    try:
        proc = await asyncio.create_subprocess_exec(
            sys.executable, str(script),
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            cwd=str(tmp_dir))
        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=30)
        out = stdout.decode(errors="replace").strip()
        err = stderr.decode(errors="replace").strip()
        return {"status": "success" if proc.returncode == 0 else "error",
                "stdout": out[:10000], "stderr": err[:5000],
                "returncode": proc.returncode}
    except asyncio.TimeoutError:
        return {"status": "error", "error": "Execution timed out after 30s"}
    except Exception as e:
        return {"status": "error", "error": str(e)}
    finally:
        if script.exists():
            script.unlink()

async def tool_read_file(path: str) -> dict:
    p = Path(path).expanduser().resolve()
    if not p.exists():
        return {"status": "error", "error": f"File not found: {path}"}
    if not p.is_file():
        return {"status": "error", "error": f"Not a file: {path}"}
    try:
        async with aiofiles.open(p, "r") as f:
            content = await f.read()
        return {"status": "success", "content": content[:50000]}
    except Exception as e:
        return {"status": "error", "error": str(e)}

async def tool_write_file(path: str, content: str) -> dict:
    p = Path(path).expanduser().resolve()
    p.parent.mkdir(parents=True, exist_ok=True)
    try:
        async with aiofiles.open(p, "w") as f:
            await f.write(content)
        return {"status": "success", "path": str(p)}
    except Exception as e:
        return {"status": "error", "error": str(e)}

async def tool_list_dir(path: str) -> dict:
    p = Path(path).expanduser().resolve()
    if not p.exists():
        return {"status": "error", "error": f"Path not found: {path}"}
    if not p.is_dir():
        return {"status": "error", "error": f"Not a directory: {path}"}
    try:
        items = []
        for entry in sorted(p.iterdir()):
            items.append({"name": entry.name, "type": "dir" if entry.is_dir() else "file",
                          "size": entry.stat().st_size if entry.is_file() else 0})
        return {"status": "success", "path": str(p), "items": items}
    except Exception as e:
        return {"status": "error", "error": str(e)}

async def tool_web_search(query: str) -> dict:
    try:
        loop = asyncio.get_event_loop()
        results = await loop.run_in_executor(None, _search_ddg, query)
        return {"status": "success", "query": query, "results": results}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def _search_ddg(query: str, max_results: int = 5) -> list[dict]:
    from curl_cffi import requests
    from bs4 import BeautifulSoup
    r = requests.get(
        "https://lite.duckduckgo.com/lite/",
        params={"q": query, "kp": "-2"},
        impersonate="chrome",
        timeout=15
    )
    soup = BeautifulSoup(r.text, "lxml")
    rows = soup.select("tr")
    out = []
    i = 0
    while i < len(rows) and len(out) < max_results:
        a = rows[i].select_one("a.result-link")
        if not a:
            i += 1
            continue
        title = a.get_text(strip=True)
        raw_url = a.get("href", "")
        snippet = ""
        if i + 1 < len(rows):
            snip_td = rows[i + 1].select_one("td.result-snippet")
            if snip_td:
                snippet = snip_td.get_text(strip=True)
        out.append({"title": title[:200], "snippet": snippet[:500], "url": raw_url[:300]})
        i += 3
    if not out:
        out.append({"title": "No results", "snippet": f"No search results for: {query}", "url": ""})
    return out

async def tool_system_now() -> dict:
    from datetime import datetime
    return {"status": "success", "datetime": datetime.now().isoformat()}

TOOL_IMPLS = {
    "system.shell.exec": tool_shell_exec,
    "system.code_interpreter": tool_code_interpreter,
    "agent.workspace.read_file": tool_read_file,
    "agent.workspace.write_file": tool_write_file,
    "agent.workspace.list_dir": tool_list_dir,
    "web.search": tool_web_search,
    "system.now": tool_system_now,
}

async def call_llm(messages: list[dict], tools: list[dict] = None,
                   tool_choice: str = "auto", model: str = None,
                   max_tokens: int = 4096, temperature: float = 0.7) -> dict:
    body = {"messages": messages, "max_tokens": max_tokens, "temperature": temperature}
    if model:
        body["model"] = model
    if tools:
        body["tools"] = tools
        body["tool_choice"] = tool_choice
    async with aiohttp.ClientSession() as sess:
        async with sess.post(f"{LLAMA_CHAT}/v1/chat/completions",
                             json=body, timeout=aiohttp.ClientTimeout(total=120)) as resp:
            if resp.status != 200:
                txt = await resp.text()
                raise HTTPException(status_code=resp.status, detail=txt[:500])
            return await resp.json()

async def stream_llm(messages: list[dict], tools: list[dict] = None,
                     tool_choice: str = "auto", model: str = None):
    body = {"messages": messages, "max_tokens": 4096, "temperature": 0.7,
            "stream": True}
    if model:
        body["model"] = model
    if tools:
        body["tools"] = tools
        body["tool_choice"] = tool_choice
    async with aiohttp.ClientSession() as sess:
        async with sess.post(f"{LLAMA_CHAT}/v1/chat/completions",
                             json=body, timeout=aiohttp.ClientTimeout(total=120)) as resp:
            if resp.status != 200:
                yield f"data: {json.dumps({'type': 'error', 'detail': await resp.text()})}\n\n"
                return
            buffer = ""
            async for chunk in resp.content.iter_chunks():
                data = chunk[0].decode(errors="replace")
                buffer += data
                while "\n" in buffer:
                    line, buffer = buffer.split("\n", 1)
                    line = line.strip()
                    if not line:
                        continue
                    if line.startswith("data: "):
                        payload = line[6:]
                        if payload == "[DONE]":
                            yield f"data: {json.dumps({'type': 'done'})}\n\n"
                            return
                        try:
                            obj = json.loads(payload)
                            yield f"data: {json.dumps({'type': 'token', 'data': obj})}\n\n"
                        except json.JSONDecodeError:
                            pass

def get_thread(thread_id: str) -> list[dict]:
    if thread_id not in conversations:
        conversations[thread_id] = [
            {"role": "system", "content": textwrap.dedent("""\
                You are an AI assistant with full system access.
                Use tools to answer questions requiring data or actions.
                Never describe what tools you would use — just use them.
                For simple greetings or chat, reply directly without tools.
                After your final answer, suggest 3 brief follow-up questions
                the user might want to ask next. Wrap them as a JSON object
                on its own line: {"suggestions": ["q1?", "q2?", "q3?"]}.
            """).strip()}
        ]
    return conversations[thread_id]

def trim_thread(thread: list[dict]):
    if len(thread) > MAX_HISTORY:
        system = [m for m in thread if m["role"] == "system"]
        rest = [m for m in thread if m["role"] != "system"]
        thread[:] = system + rest[-(MAX_HISTORY - len(system)):]

@app.get("/health")
async def health():
    async with aiohttp.ClientSession() as sess:
        try:
            async with sess.get(f"{LLAMA_CHAT}/health",
                                timeout=aiohttp.ClientTimeout(total=3)) as r:
                chat = r.status == 200
        except:
            chat = False
        try:
            async with sess.get(f"{LLAMA_EMBED}/health",
                                timeout=aiohttp.ClientTimeout(total=3)) as r:
                embed = r.status == 200
        except:
            embed = False
    return {"status": "ok", "chat": chat, "embeddings": embed, "tools": list(TOOL_IMPLS.keys())}

@app.get("/", response_class=HTMLResponse)
async def index():
    p = Path(__file__).parent / "templates" / "index.html"
    return HTMLResponse(p.read_text())

@app.post("/v1/chat")
async def chat(request: Request):
    body = await request.json()
    msg = body.get("message", "").strip()
    thread_id = body.get("thread_id", "default")
    if not msg:
        raise HTTPException(400, "message is required")
    thread = get_thread(thread_id)
    thread.append({"role": "user", "content": msg})
    trim_thread(thread)
    return await _process_tools(thread, thread_id)

@app.post("/v1/chat/stream")
async def chat_stream(request: Request):
    body = await request.json()
    thread_id = body.get("thread_id", "default")
    proposal_id = body.get("proposal_id")
    decisions = body.get("decisions")
    if proposal_id and decisions is not None:
        return StreamingResponse(_execute_proposal_stream(proposal_id, decisions, thread_id),
                                 media_type="text/event-stream")
    msg = body.get("message", "").strip()
    if not msg:
        raise HTTPException(400, "message is required")
    thread = get_thread(thread_id)
    thread.append({"role": "user", "content": msg})
    trim_thread(thread)
    return StreamingResponse(_process_stream(thread, thread_id),
                             media_type="text/event-stream")

@app.delete("/v1/conversation/{thread_id}")
async def clear_conversation(thread_id: str):
    conversations.pop(thread_id, None)
    return {"status": "ok"}

@app.get("/v1/tools")
async def list_tools():
    return TOOL_DEFINITIONS

@app.post("/v1/chat/propose")
async def chat_propose(request: Request):
    body = await request.json()
    msg = body.get("message", "").strip()
    thread_id = body.get("thread_id", "default")
    if not msg:
        raise HTTPException(400, "message is required")
    thread = get_thread(thread_id)
    temp_msgs = thread + [{"role": "user", "content": msg}]
    resp = await call_llm(temp_msgs, tools=TOOL_DEFINITIONS, tool_choice="auto")
    msg_obj = resp["choices"][0]["message"]
    tool_calls = msg_obj.get("tool_calls")
    if not tool_calls:
        thread.append({"role": "user", "content": msg})
        content = msg_obj.get("content", "") or ""
        thread.append({"role": "assistant", "content": content})
        trim_thread(thread)
        return {"response": content, "thread_id": thread_id, "tool_calls": []}
    pid = f"{thread_id}_{len(thread)}_{int(asyncio.get_event_loop().time()*1000000)}"
    raw = []
    for tc in tool_calls:
        try:
            a = json.loads(tc["function"]["arguments"])
        except:
            a = {}
        raw.append({"tool_call_id": tc["id"], "name": tc["function"]["name"], "args": a})
    proposals[pid] = {
        "thread_id": thread_id, "user_content": msg,
        "assistant_content": msg_obj.get("content") or "",
        "tool_calls": tool_calls,
    }
    return {"proposal_id": pid, "tool_calls": raw, "content": msg_obj.get("content") or "",
            "thread_id": thread_id}

async def _process_tools(thread: list[dict], thread_id: str, depth: int = 0):
    if depth > 5:
        return {"response": "Tool call depth exceeded (max 5).", "thread_id": thread_id}
    resp = await call_llm(thread, tools=TOOL_DEFINITIONS, tool_choice="auto")
    msg = resp["choices"][0]["message"]
    tool_calls = msg.get("tool_calls")
    if tool_calls:
        thread.append({"role": "assistant", "content": msg.get("content") or "",
                       "tool_calls": tool_calls})
        results = []
        for tc in tool_calls:
            fn = tc["function"]
            name = fn["name"]
            try:
                args = json.loads(fn["arguments"])
            except:
                args = {}
            impl = TOOL_IMPLS.get(name)
            if impl:
                result = await impl(**args)
            else:
                result = {"status": "error", "error": f"Unknown tool: {name}"}
            results.append({"tool_call_id": tc["id"], "name": name, "result": result})
            thread.append({"role": "tool", "tool_call_id": tc["id"],
                          "content": json.dumps(result)})
        trim_thread(thread)
        final = await _process_tools(thread, thread_id, depth + 1)
        final["tool_results"] = results
        return final
    else:
        content = msg.get("content", "") or ""
        thread.append({"role": "assistant", "content": content})
        trim_thread(thread)
        return {"response": content, "thread_id": thread_id}

async def _process_stream(thread: list[dict], thread_id: str, depth: int = 0):
    if depth > 5:
        yield f"data: {json.dumps({'type': 'error', 'detail': 'Tool call depth exceeded (max 5)'})}\n\n"
        return
    collected_content = ""
    collected_tool_calls = None
    async for event in stream_llm(thread, tools=TOOL_DEFINITIONS, tool_choice="auto"):
        yield event
        try:
            data = json.loads(event[6:].strip()) if event.startswith("data: ") else None
        except:
            data = None
        if data is None:
            continue
        if data.get("type") == "done":
            break
        obj = data.get("data", {})
        delta = obj.get("choices", [{}])[0].get("delta", {})
        collected_content += delta.get("content", "") or ""
        if "tool_calls" in delta:
            if collected_tool_calls is None:
                collected_tool_calls = []
            for tc in delta["tool_calls"]:
                idx = tc.get("index", len(collected_tool_calls))
                while len(collected_tool_calls) <= idx:
                    collected_tool_calls.append({"id": "", "function": {"name": "", "arguments": ""}})
                if tc.get("type"):
                    collected_tool_calls[idx]["type"] = tc["type"]
                if tc.get("id"):
                    collected_tool_calls[idx]["id"] = tc["id"]
                if tc.get("function", {}).get("name"):
                    collected_tool_calls[idx]["function"]["name"] += tc["function"]["name"]
                if tc.get("function", {}).get("arguments"):
                    collected_tool_calls[idx]["function"]["arguments"] += tc["function"]["arguments"]
    if collected_tool_calls:
        msg = {"role": "assistant", "content": collected_content,
               "tool_calls": collected_tool_calls}
        thread.append(msg)
        results = []
        for tc in collected_tool_calls:
            fn = tc["function"]
            name = fn["name"]
            try:
                args = json.loads(fn["arguments"])
            except:
                args = {}
            impl = TOOL_IMPLS.get(name)
            yield f"data: {json.dumps({'type': 'tool_start', 'name': name, 'args': args})}\n\n"
            if impl:
                result = await impl(**args)
            else:
                result = {"status": "error", "error": f"Unknown tool: {name}"}
            results.append({"tool_call_id": tc["id"], "name": name, "result": result})
            yield f"data: {json.dumps({'type': 'tool_result', 'name': name, 'result': result})}\n\n"
            thread.append({"role": "tool", "tool_call_id": tc["id"],
                          "content": json.dumps(result)})
        trim_thread(thread)
        async for event in _process_stream(thread, thread_id, depth + 1):
            yield event
    else:
        if collected_content:
            display, suggestions = _extract_suggestions(collected_content)
            if display != collected_content:
                yield f"data: {json.dumps({'type': 'suggestions', 'suggestions': suggestions})}\n\n"
            thread.append({"role": "assistant", "content": collected_content})
            trim_thread(thread)

def _extract_suggestions(text: str) -> tuple[str, list[str]]:
    m = re.search(r'\{\s*"suggestions"\s*:\s*\[(.*?)\]\s*\}', text, re.DOTALL)
    if not m:
        return text, []
    try:
        obj = json.loads("{" + m.group(0) + "}")
        suggestions = obj.get("suggestions", [])[:3]
        rest = text[:m.start()].rstrip() + text[m.end():]
        return rest, suggestions
    except:
        return text, []

async def _execute_proposal_stream(proposal_id: str, decisions: list[dict], thread_id: str):
    if proposal_id not in proposals:
        yield f"data: {json.dumps({'type': 'error', 'detail': 'Proposal not found'})}\n\n"
        return
    prop = proposals.pop(proposal_id)
    thread = get_thread(thread_id)
    thread.append({"role": "user", "content": prop["user_content"]})
    tool_calls = prop["tool_calls"]
    assistant_content = prop["assistant_content"]
    thread.append({"role": "assistant", "content": assistant_content, "tool_calls": tool_calls})
    decision_map = {d["tool_call_id"]: d.get("approved", False) for d in decisions}
    for tc in tool_calls:
        tc_id = tc["id"]
        fn = tc["function"]
        name = fn["name"]
        try:
            args = json.loads(fn["arguments"]) if isinstance(fn["arguments"], str) else fn["arguments"]
        except:
            args = {}
        yield f"data: {json.dumps({'type': 'tool_start', 'name': name, 'args': args, 'tool_call_id': tc_id})}\n\n"
        if decision_map.get(tc_id, False):
            impl = TOOL_IMPLS.get(name)
            if impl:
                result = await impl(**args)
            else:
                result = {"status": "error", "error": f"Unknown tool: {name}"}
        else:
            result = {"status": "skipped", "reason": "Rejected by user"}
        yield f"data: {json.dumps({'type': 'tool_result', 'name': name, 'result': result, 'tool_call_id': tc_id})}\n\n"
        thread.append({"role": "tool", "tool_call_id": tc_id, "content": json.dumps(result)})
    trim_thread(thread)
    collected_content = ""
    async for event in stream_llm(thread, tools=None):
        yield event
        try:
            data = json.loads(event[6:].strip()) if event.startswith("data: ") else None
        except:
            data = None
        if data and data.get("type") == "done":
            pass
        elif data and data.get("type") == "token":
            delta = data.get("data", {}).get("choices", [{}])[0].get("delta", {})
            collected_content += delta.get("content", "") or ""
    if collected_content:
        display, suggestions = _extract_suggestions(collected_content)
        if display != collected_content:
            yield f"data: {json.dumps({'type': 'suggestions', 'suggestions': suggestions})}\n\n"
        thread.append({"role": "assistant", "content": collected_content})
        trim_thread(thread)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080)