#!/usr/bin/env python3
"""
SHIRE Brain API Server.

Laptop-only bridge between ARMOR Pi Core and local Ollama.
Binds to the laptop Tailscale IP so no public router port is needed.
"""

import json
import os
import socket
import subprocess
import sys
import time
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path


ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from codex.prompt_builder import build_prompt
from services.blender_brain_contract import (
    BlenderBrainContractError,
    DEEP_MODEL as BLENDER_DEEP_MODEL,
    FAST_MODEL as BLENDER_FAST_MODEL,
    prepare_route,
)


OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://127.0.0.1:11434").rstrip("/")
FAST_MODEL = os.environ.get("SHIRE_FAST_MODEL", "shire-fast:qwen3-1.7b")
DEEP_MODEL = os.environ.get(
    "SHIRE_DEEP_MODEL",
    os.environ.get("SHIRE_BRAIN_MODEL", "shire-brain:qwen3-8b"),
)
PORT = int(os.environ.get("SHIRE_BRAIN_PORT", "8765"))
HOST_OVERRIDE = os.environ.get("SHIRE_BRAIN_HOST", "").strip()
OLLAMA_TIMEOUT = int(os.environ.get("SHIRE_OLLAMA_TIMEOUT", "300"))


def tailscale_ip():
    try:
        result = subprocess.run(
            ["tailscale", "ip", "-4"],
            check=False,
            capture_output=True,
            text=True,
            timeout=5,
        )
        for line in result.stdout.splitlines():
            ip = line.strip()
            if ip:
                return ip
    except Exception:
        return ""
    return ""


def bind_host():
    if HOST_OVERRIDE:
        return HOST_OVERRIDE
    return tailscale_ip() or "127.0.0.1"


def json_response(handler, payload, status=200):
    body = json.dumps(payload, indent=2).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


def choose_model(prompt, mode):
    mode = (mode or "").strip().lower()
    clean_prompt = (prompt or "").strip()
    upper_prompt = clean_prompt.upper()
    lower_prompt = clean_prompt.lower()

    if mode in {"deep", "slow", "qwen8b"} or upper_prompt.startswith("DEEP "):
        if upper_prompt.startswith("DEEP "):
            clean_prompt = clean_prompt[5:].strip()
        return DEEP_MODEL, "deep", clean_prompt, "manual deep request"

    if mode in {"fast", "quick", "qwen1.7b"} or upper_prompt.startswith("FAST "):
        if upper_prompt.startswith("FAST "):
            clean_prompt = clean_prompt[5:].strip()
        return FAST_MODEL, "fast", clean_prompt, "manual fast request"

    deep_keywords = [
        "debug", "traceback", "error", "exception", "fix this", "repair",
        "architecture", "design a system", "build a system", "engineering",
        "plan", "strategy", "compare", "analyse", "analyze", "review",
        "step by step", "full guide", "deep", "complex", "hard",
        "code", "script", "patch", "refactor", "security", "threat",
        "legal", "medical", "financial", "mortgage", "ndis",
        "long answer", "detailed", "explain properly",
    ]

    if len(clean_prompt) > 220:
        return DEEP_MODEL, "deep", clean_prompt, "long prompt"

    for keyword in deep_keywords:
        if keyword in lower_prompt:
            return DEEP_MODEL, "deep", clean_prompt, f"matched keyword: {keyword}"

    return FAST_MODEL, "fast", clean_prompt, "simple prompt"


def blender_system_context(route):
    """Build an isolated, non-executing Blender specialist context."""
    denied = sorted(
        name
        for name, allowed in route["capabilities"].items()
        if allowed is False
    )

    return (
        "SHIRE BLENDER SPECIALIST ROUTE — DRAFT ONLY\n"
        f"Task: {route['task']}\n"
        f"Project: {route['project_id']}\n"
        f"Revision: {route['revision']}\n"
        f"Bound artifact SHA-256: {route['artifact_sha256']}\n"
        f"Authority effect: {route['authority_effect']}\n"
        "Denied capabilities: " + ", ".join(denied) + "\n"
        "Use only the supplied prompt and binding. "
        "Do not call tools, access files, browse, run commands, run scripts, "
        "invoke Blender, invoke BlenderMCP, render, export, install models, "
        "modify project state, grant approval, or claim execution occurred. "
        "Return only the requested draft content."
    )


def resolve_request_route(payload, prompt):
    """Resolve locked Blender routing or preserve normal SHIRE routing."""
    blender_task = str(
        payload.get("blender_task", "")
    ).strip().lower()

    if blender_task:
        route = prepare_route(
            blender_task,
            payload.get("project_id", ""),
            payload.get("revision"),
            payload.get("artifact_sha256", ""),
        )

        model = route["model"]

        if model == BLENDER_FAST_MODEL:
            mode = "blender_fast"
        elif model == BLENDER_DEEP_MODEL:
            mode = "blender_deep"
        else:
            raise BlenderBrainContractError(
                "Blender route selected an unapproved model"
            )

        return (
            model,
            mode,
            str(prompt).strip(),
            f"locked Blender task: {blender_task}",
            route["max_tokens"],
            route,
        )

    requested_tokens = int(payload.get("max_tokens", 180))
    max_tokens = max(1, min(requested_tokens, 800))

    model, mode, clean_prompt, route_reason = choose_model(
        prompt,
        payload.get("mode", ""),
    )

    return (
        model,
        mode,
        clean_prompt,
        route_reason,
        max_tokens,
        None,
    )



def guarded_prompt(prompt, mode):
    clean_prompt = (prompt or "").strip()

    if mode != "deep":
        return clean_prompt

    guard = """
ARMOR DEEP BRAIN SAFETY GUARD:
- If the user asks for debugging or repair but gives no real traceback/log/output/file content, ask for the missing output first.
- Do not invent service names, file paths, package names, commands, logs, or results.
- Do not suggest sudo install/delete/restart commands unless the supplied evidence supports them.
- Prefer one safe diagnostic step.
- Keep it short.
"""

    return guard.strip() + "\n\nUSER REQUEST:\n" + clean_prompt

def ask_ollama(model, prompt, max_tokens, system_context=""):
    messages = []
    if system_context:
        messages.append({"role": "system", "content": system_context})
    messages.append({"role": "user", "content": prompt})

    payload = {
        "model": model,
        "think": False,
        "stream": False,
        "messages": messages,
        "options": {
            "temperature": 0.2,
            "num_predict": max_tokens,
        },
    }

    request = urllib.request.Request(
        OLLAMA_URL + "/api/chat",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    with urllib.request.urlopen(request, timeout=OLLAMA_TIMEOUT) as response:
        return json.loads(response.read().decode("utf-8"))


class BrainHandler(BaseHTTPRequestHandler):
    server_version = "SHIREBrainAPI/0007"

    def log_message(self, fmt, *args):
        print("%s - - [%s] %s" % (self.address_string(), self.log_date_time_string(), fmt % args))

    def do_GET(self):
        if self.path not in {"/", "/health"}:
            json_response(self, {"ok": False, "error": "Not found"}, status=404)
            return

        json_response(
            self,
            {
                "ok": True,
                "service": "SHIRE Brain API",
                "host": socket.gethostname(),
                "model": FAST_MODEL,
                "fast_model": FAST_MODEL,
                "deep_model": DEEP_MODEL,
                "blender_fast_model": BLENDER_FAST_MODEL,
                "blender_deep_model": BLENDER_DEEP_MODEL,
                "blender_routing_contract": "shire.blender_brain.route.v1",
                "knowledge_context": "ray-approved registry",
                "ollama_url": OLLAMA_URL,
                "message": "SHIRE brain online.",
            },
        )

    def do_POST(self):
        if self.path != "/ask":
            json_response(self, {"ok": False, "error": "Not found"}, status=404)
            return

        try:
            length = int(self.headers.get("Content-Length", "0"))
            raw = self.rfile.read(length).decode("utf-8")
            payload = json.loads(raw or "{}")

            prompt = str(payload.get("prompt", "")).strip()
            if not prompt:
                json_response(self, {"ok": False, "error": "Missing prompt"}, status=400)
                return

            (
                model,
                mode,
                clean_prompt,
                route_reason,
                max_tokens,
                blender_route,
            ) = resolve_request_route(payload, prompt)

            if blender_route is not None:
                guarded = clean_prompt
                system_context = blender_system_context(blender_route)
            else:
                guarded = guarded_prompt(clean_prompt, mode)
                system_context = build_prompt(clean_prompt)

            start = time.time()
            ollama_data = ask_ollama(
                model,
                guarded,
                max_tokens,
                system_context=system_context,
            )
            elapsed = round(time.time() - start, 2)

            message = ollama_data.get("message", {})
            answer = message.get("content", "").strip()

            json_response(
                self,
                {
                    "ok": True,
                    "mode": mode,
                    "model": model,
                    "route_reason": route_reason,
                    "blender_route": blender_route,
                    "elapsed_seconds": elapsed,
                    "answer": answer,
                },
            )

        except BlenderBrainContractError as exc:
            json_response(
                self,
                {
                    "ok": False,
                    "error": f"Blender routing blocked: {exc}",
                },
                status=400,
            )
        except urllib.error.HTTPError as exc:
            try:
                detail = exc.read().decode("utf-8", errors="replace")
            except Exception:
                detail = str(exc)
            json_response(
                self,
                {
                    "ok": False,
                    "error": f"Ollama HTTP error: {exc.code}",
                    "detail": detail,
                },
                status=500,
            )
        except Exception as exc:
            json_response(
                self,
                {
                    "ok": False,
                    "error": str(exc),
                },
                status=500,
            )


def main():
    host = bind_host()
    server = ThreadingHTTPServer((host, PORT), BrainHandler)
    print(f"SHIRE Brain API listening on http://{host}:{PORT}")
    print(f"Fast model: {FAST_MODEL}")
    print(f"Deep model: {DEEP_MODEL}")
    print(f"Ollama URL: {OLLAMA_URL}")
    server.serve_forever()


if __name__ == "__main__":
    main()
