#!/usr/bin/env python3
"""
pimage-daemon.py — Watches for image processing requests, sends them to
a vision model API (OpenAI-compatible), and returns results.

Protocol (dirs default to ~/.pimage, override with PIMAGE_DIR):
  Signal:  <PIMAGE_DIR>/current               — contains UUID of the current request
  Request: <PIMAGE_DIR>/in/<uuid>/request.json — { id, image_path, instruction }
  Image:   <PIMAGE_DIR>/in/<uuid>/image         — copy of the image file
  Result:  <PIMAGE_DIR>/out/<uuid>/result.json  — { success, content, error, cost }

Usage:
  python3 pimage-daemon.py [--once]

Env vars:
  PIMAGE_DIR        state dir (default ~/.pimage)
  PIMAGE_MODEL      model (default opencode-go/mimo-v2.5)
  PIMAGE_API_BASE   API base URL (default https://opencode.ai/zen/go/v1)
  OPENCODE_API_KEY  API key (required)
"""

import argparse
import base64
import json
import mimetypes
import os
import re
import signal
import sys
import time
from pathlib import Path
from urllib.request import Request, urlopen
from urllib.error import URLError

# Config — all paths env-overridable (no hardcoded user paths)
PIDIR = Path(os.environ.get("PIMAGE_DIR", str(Path.home() / ".pimage")))
INBOX = PIDIR / "in"
OUTBOX = PIDIR / "out"
SIGNAL = PIDIR / "current"
POLL_INTERVAL = float(os.environ.get("PIMAGE_POLL_INTERVAL", "0.3"))
API_TIMEOUT = int(os.environ.get("PIMAGE_API_TIMEOUT", "120"))


def log(msg: str):
    ts = time.strftime("%H:%M:%S")
    print(f"[pimage] {ts} {msg}", flush=True)


def base64_encode(path: Path) -> tuple[str, str]:
    data = path.read_bytes()
    encoded = base64.b64encode(data).decode("ascii")
    mime, _ = mimetypes.guess_type(str(path))
    if mime is None:
        ext = path.suffix.lower()
        mime = {
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".gif": "image/gif",
            ".webp": "image/webp",
            ".bmp": "image/bmp",
        }.get(ext, "image/png")
    return encoded, mime


class VisionAPI:
    """Calls the opencode-go API directly with images."""

    def __init__(self, model: str, api_key: str, base_url: str | None = None):
        self.base_url = base_url or os.environ.get("PIMAGE_API_BASE", "https://opencode.ai/zen/go/v1")
        self.model = model
        self.api_key = api_key
        self.base_url = self.base_url.rstrip("/")

    def send_prompt(self, message: str, image_path: str | Path | None = None) -> tuple[str | None, dict | None]:
        """
        Send a prompt (optionally with an image) to the vision model API.
        Returns (text_response, cost_info) or (None, None) on failure.
        """
        content: list[dict] = [{"type": "text", "text": message}]

        if image_path and os.path.exists(image_path):
            img_path = Path(image_path)
            try:
                b64_data, mime_type = base64_encode(img_path)
                content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:{mime_type};base64,{b64_data}"},
                })
                log(f"  Image: {img_path.name} ({mime_type}, {len(b64_data)} bytes base64)")
            except Exception as e:
                log(f"  WARN: failed to encode image {img_path}: {e}")

        messages: list = []
        has_image = image_path is not None and os.path.exists(image_path)
        if has_image:
            messages.append({
                "role": "system",
                "content": (
                    "You are a precise vision analysis model. Describe images with exhaustive detail. "
                    "Leave nothing out. Transcribe EVERY visible text character for character. "
                    "Cover layout, panels, sections, hierarchy, indentation, UI elements (buttons, icons, "
                    "inputs, scrollbars), highlights/borders/shadows, color scheme, selected/hovered states, "
                    "collapsed/expanded sections, and application context. Be exhaustive do not summarize."
                ),
            })

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

        payload = {
            "model": self.model,
            "messages": messages,
            "stream": False,
            "max_tokens": 4096,
        }

        payload_json = json.dumps(payload).encode("utf-8")
        log(f"  Sending request ({len(message)} chars, {len(payload_json)} bytes)...")

        url = f"{self.base_url}/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
            "User-Agent": "pimage/1.0",
            "Accept": "application/json",
        }

        try:
            req = Request(url, data=payload_json, headers=headers, method="POST")
            with urlopen(req, timeout=API_TIMEOUT) as response:
                resp_data = json.loads(response.read().decode("utf-8"))

            choices = resp_data.get("choices", [])
            if not choices:
                log(f"  ERROR: no choices in response")
                return None, None

            message = choices[0].get("message", {})
            content_text = message.get("content", "")

            # Extract usage/cost
            usage = resp_data.get("usage", {})
            cost_raw = resp_data.get("cost", "0")
            cost_value = float(cost_raw) if isinstance(cost_raw, str) else (cost_raw if isinstance(cost_raw, (int, float)) else 0)
            cost_info = {
                "model": self.model,
                "tokens": usage.get("total_tokens", 0),
                "prompt_tokens": usage.get("prompt_tokens", 0),
                "completion_tokens": usage.get("completion_tokens", 0),
                "cost": cost_value,
            }

            # Strip thinking blocks if present
            if content_text and "<thinking>" in content_text:
                content_text = re.sub(r'<thinking>.*?</thinking>', '', content_text, flags=re.DOTALL).strip()

            if not content_text:
                finish_reason = choices[0].get("finish_reason", "unknown")
                log(f"  WARN: empty response, finish: {finish_reason}")
                return "(no text response)", cost_info

            log(f"  Response: {len(content_text)} chars, cost: ${cost_value:.6f}")
            return content_text, cost_info

        except URLError as e:
            log(f"  ERROR: API request failed: {e}")
            if hasattr(e, "read"):
                try:
                    error_body = e.read().decode("utf-8", errors="replace")
                    log(f"  Error body: {error_body[:500]}")
                except Exception:
                    pass
            return None, None
        except json.JSONDecodeError as e:
            log(f"  ERROR: failed to parse response: {e}")
            return None, None
        except Exception as e:
            log(f"  ERROR: unexpected error: {e}")
            return None, None


def process_request(api: VisionAPI, uuid: str):
    req_file = INBOX / uuid / "request.json"
    img_file = INBOX / uuid / "image"
    out_dir = OUTBOX / uuid
    result_file = out_dir / "result.json"

    if not req_file.exists():
        log(f"WARN: {uuid} has no request.json, skipping")
        _clear_signal()
        return

    out_dir.mkdir(parents=True, exist_ok=True)

    try:
        with open(req_file) as f:
            req = json.load(f)
    except Exception as e:
        log(f"ERROR: failed to parse {req_file}: {e}")
        _write_result(result_file, success=False, error=f"Invalid request: {e}")
        _clear_signal()
        return

    instruction = req.get("instruction") or req.get("prompt") or "Describe this image in detail"
    image_path = str(img_file) if img_file.exists() else req.get("image_path")

    log(f"-> Processing {uuid}")
    log(f"   Instruction: {instruction[:100]}...")

    response, cost_info = api.send_prompt(instruction, image_path=image_path)

    if response is None:
        log(f"X Failed for {uuid}")
        _write_result(result_file, success=False, error="Vision API returned no response")
    else:
        log(f"OK Done: {uuid} ({len(response)} chars)")
        _write_result(result_file, success=True, content=response, cost=cost_info)

    _clear_signal()


def _write_result(path: Path, success: bool, content: str = "", error: str = "", cost: dict | None = None):
    result = {
        "success": success,
        "content": content,
        "error": error,
    }
    if cost:
        result["cost"] = cost
    path.write_text(json.dumps(result, indent=2))


def _clear_signal():
    try:
        SIGNAL.unlink(missing_ok=True)
    except Exception:
        pass


def _resolve_model(model_spec: str) -> str:
    if "/" in model_spec:
        return model_spec.split("/", 1)[1]
    return model_spec


def main():
    parser = argparse.ArgumentParser(description="pimage vision agent daemon")
    parser.add_argument("--once", action="store_true", help="Process one request and exit")
    parser.add_argument(
        "--model",
        default=os.environ.get("PIMAGE_MODEL", "opencode-go/mimo-v2.5"),
        help="Vision model (default: opencode-go/mimo-v2.5)",
    )
    args = parser.parse_args()

    model_id = _resolve_model(args.model)
    api_key = os.environ.get("OPENCODE_API_KEY")
    if not api_key:
        log("FATAL: OPENCODE_API_KEY environment variable not set")
        sys.exit(1)

    INBOX.mkdir(parents=True, exist_ok=True)
    OUTBOX.mkdir(parents=True, exist_ok=True)

    log("pimage daemon starting")
    log(f"   Model: {model_id}")
    log(f"   Watch: {SIGNAL}")
    log(f"   Mode:  {'single-shot' if args.once else 'continuous'}")
    print("---", flush=True)

    api = VisionAPI(model=model_id, api_key=api_key)

    def shutdown(signum, frame):
        log("Shutting down...")
        sys.exit(0)

    signal.signal(signal.SIGINT, shutdown)
    signal.signal(signal.SIGTERM, shutdown)

    try:
        if args.once:
            if SIGNAL.exists():
                uuid = SIGNAL.read_text().strip()
                process_request(api, uuid)
        else:
            log(f"Watching {SIGNAL} (poll {POLL_INTERVAL}s)...")
            while True:
                if SIGNAL.exists():
                    uuid = SIGNAL.read_text().strip()
                    if uuid:
                        process_request(api, uuid)
                time.sleep(POLL_INTERVAL)
    except Exception as e:
        log(f"FATAL: {e}")
        raise


if __name__ == "__main__":
    main()