"""
Analyzer MCP Server (FastMCP)
- Tool: classify_pii(image_path: str) -> {"instances": [...]}
- Reads the system prompt from prompts/pii_system_prompt.md (for reference)

Vision Model Provider Selection:
- Set USE_GEMINI_VISION=true to use Google Gemini (paid, cloud)
- Set USE_OLLAMA_VISION=true to use Ollama with Qwen2.5-VL (free, local, private)
- Set USE_OLLAMA_VISION=false and USE_GEMINI_VISION=false to use OpenAI Vision API (paid, cloud)

Gemini Vision Configuration (when USE_GEMINI_VISION=true):
- Set GEMINI_VISION_MODEL to choose the model (default: gemini-2.0-flash-exp)
- Set GEMINI_VISION_SEED for deterministic outputs (optional)
- Requires GEMINI_API_KEY environment variable

Ollama Vision Configuration (when USE_OLLAMA_VISION=true):
- Set OLLAMA_VISION_MODEL to choose the model (default: qwen2.5-vl:32b)
- Requires Ollama installed with vision model: ollama pull qwen2.5-vl:32b
- Set OLLAMA_BASE_URL if not using default (http://localhost:11434)
- Set OLLAMA_VISION_SEED for deterministic outputs (optional)

OpenAI Configuration (when USE_OLLAMA_VISION=false and USE_GEMINI_VISION=false):
- Set OPENAI_ANALYZER_MODEL to choose the OpenAI model (default: gpt-4o-mini)
- Falls back to OPENAI_VISION_MODEL, then to 'gpt-4o-mini' if unset
- Set OPENAI_ANALYZER_SEED (or OPENAI_SEED) for deterministic outputs (optional)
- Requires OPENAI_API_KEY environment variable

This is a runnable MCP server over stdio, modeled after the reference servers.
Note: Tool handlers here are synchronous and should not create unawaited coroutines.
"""

from __future__ import annotations

from pathlib import Path
import sys
import os
import base64
import mimetypes
import json
from typing import Any, Dict, Union

try:
    from PIL import Image
except ImportError:
    Image = None  # type: ignore

# Check PIL availability early
if Image is None:
    print("[ERROR] PIL (Pillow) not available. Install with: pip install Pillow", file=sys.stderr, flush=True)
    sys.exit(1)


try:
    from dotenv import load_dotenv  # type: ignore
except Exception:  # pragma: no cover
    load_dotenv = None

# Load .env early if available
if load_dotenv:
    load_dotenv(dotenv_path=".env")

# Check which vision provider to use
use_gemini_vision = os.getenv("USE_GEMINI_VISION", "false").lower() in ("true", "1", "yes")
use_ollama_vision = os.getenv("USE_OLLAMA_VISION", "true").lower() in ("true", "1", "yes")


def _project_root() -> Path:
    # mcp/servers -> mcp -> repo root
    return Path(__file__).resolve().parents[2]


def _load_prompt() -> str:
    prompt_path = _project_root() / "prompts" / "pii_system_prompt.md"
    if prompt_path.exists():
        return prompt_path.read_text(encoding="utf-8")
    return ""


def _encode_image_as_data_url(image_input: Any) -> str:
    if Image and isinstance(image_input, Image.Image):
        # PIL Image - encode directly
        import io
        buf = io.BytesIO()
        image_input.save(buf, format="JPEG", quality=90)
        b = buf.getvalue()
        mime = "image/jpeg"
    else:
        # Path - read from file
        mime, _ = mimetypes.guess_type(str(image_input))
        if not mime:
            # default to jpeg if unknown
            mime = "image/jpeg"
        b = image_input.read_bytes()
    
    b64 = base64.b64encode(b).decode("ascii")
    return f"data:{mime};base64,{b64}"


def _run_without_mcp() -> None:
    print(
        "Analyzer MCP server requires the 'mcp' package.\n"
        "Install with: pip install mcp  # or your chosen MCP SDK\n"
        "Then run: python -m mcp.servers.analyzer_server"
    )


try:
    from mcp.server.fastmcp import FastMCP
except Exception:
    FastMCP = None  # type: ignore[misc,assignment]


def _classify_impl(image_path: str) -> Dict[str, Any]:
    """Shared implementation used by both MCP tool and direct import."""
    ipath = Path(image_path)
    if not ipath.exists():
        return {"instances": [], "error": f"image not found: {image_path}"}

    # Get image dimensions for prompt
    try:
        with Image.open(ipath) as img:
            img_width, img_height = img.size
    except Exception as e:
        return {"instances": [], "error": f"failed to read image dimensions: {e}"}

    # Build instructions/system prompt from file
    base_prompt = _load_prompt()
    if not base_prompt:
        print(f"[DEBUG] ERROR: no prompt", file=sys.stderr, flush=True)
    
    # Add image dimensions to the prompt
    system_msg = (
        f"\n\nIMAGE DIMENSIONS: The image has dimensions {img_width} x {img_height} pixels (width x height). "
        "Use this information for calculating bounding box coordinates and sizes."
        + base_prompt.strip()
    )

    # Prepare image data as data URL for chat.completions
    try:
        data_url = _encode_image_as_data_url(ipath)
    except Exception as e:  # pragma: no cover
        return {"instances": [], "error": f"failed to read image: {e}"}

    # Check which provider to use
    use_gemini = os.getenv("USE_GEMINI_VISION", "false").lower() in ("true", "1", "yes")
    use_ollama = os.getenv("USE_OLLAMA_VISION", "true").lower() in ("true", "1", "yes")
    
    if use_gemini:
        # Google Gemini Vision - paid, cloud
        # Uses OpenAI-compatible endpoint for simplicity
        model = os.getenv("GEMINI_VISION_MODEL", "gemini-2.0-flash-exp")
        base_url = "https://generativelanguage.googleapis.com/v1beta/openai/"
        seed_env = os.getenv("GEMINI_VISION_SEED")
        provider = "gemini"
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            return {
                "instances": [],
                "error": "GEMINI_API_KEY env var not set; cannot run Gemini Analyzer",
            }
    elif use_ollama:
        # Ollama Vision (Qwen2.5-VL) - free, local
        model = os.getenv("OLLAMA_VISION_MODEL", "qwen2.5vl:32b")
        base_url = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1")
        seed_env = os.getenv("OLLAMA_VISION_SEED")
        provider = "ollama"
        api_key = "ollama"  # Ollama doesn't require real API key
    else:
        # OpenAI Vision API - paid, cloud
        model = os.getenv("OPENAI_ANALYZER_MODEL") or os.getenv("OPENAI_VISION_MODEL") or "gpt-4o-mini"
        base_url = None  # Use default OpenAI endpoint
        seed_env = os.getenv("OPENAI_ANALYZER_SEED") or os.getenv("OPENAI_SEED")
        provider = "openai"
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            return {
                "instances": [],
                "error": "OPENAI_API_KEY env var not set; cannot run OpenAI Analyzer",
            }

    # Optional deterministic seed
    seed: int | None
    try:
        seed = int(seed_env) if seed_env is not None else None
        if seed is not None and seed < 0:
            seed = None
    except Exception:
        seed = None
    seed_applied = False

    try:
        # Helper to call the model on a given data_url and clean the result
        def _classify_data_url(msg_system: str, data_url_in: str) -> tuple[list[dict[str, Any]], int]:
            nonlocal seed_applied

            # Use appropriate API for each provider
            if provider == "gemini":
                # Google Gemini via OpenAI-compatible endpoint
                try:
                    from openai import OpenAI  # type: ignore
                except Exception:
                    return [], 1
                
                client = OpenAI(api_key=api_key, base_url=base_url)
                
                common_args = {
                    "model": model,
                    "temperature": 0,
                    "response_format": {"type": "json_object"},
                }
                messages1 = [
                    {"role": "system", "content": msg_system},
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": data_url_in, "detail": "high"}},
                        ],
                    },
                ]
                try:
                    comp = (
                        client.chat.completions.create(messages=messages1, seed=seed, **common_args)
                        if seed is not None
                        else client.chat.completions.create(messages=messages1, **common_args)
                    )
                    if seed is not None:
                        seed_applied = True
                except TypeError:
                    comp = client.chat.completions.create(messages=messages1, **common_args)

                content_loc = comp.choices[0].message.content or "{}"
                result_loc = json.loads(content_loc)
                
            elif provider == "ollama":
                # Native Ollama API (/api/generate) - works better than OpenAI-compatible endpoint
                import requests
                
                # Extract base64 image from data URL
                if "base64," in data_url_in:
                    img_b64 = data_url_in.split("base64,")[1]
                else:
                    img_b64 = data_url_in
                
                # Remove /v1 from base_url if present
                ollama_url = base_url.replace("/v1", "") if base_url else "http://localhost:11434"
                
                payload = {
                    "model": model,
                    "prompt": msg_system + "\n\nAnalyze this image and return JSON.",
                    "images": [img_b64],
                    "stream": False,
                    "format": "json",
                    "options": {"temperature": 0}
                }
                if seed is not None:
                    payload["options"]["seed"] = seed
                    seed_applied = True
                
                response = requests.post(f"{ollama_url}/api/generate", json=payload)
                response.raise_for_status()
                result = response.json()
                content_loc = result.get("response", "{}")
                result_loc = json.loads(content_loc)
            else:
                # OpenAI API
                try:
                    from openai import OpenAI  # type: ignore
                except Exception:
                    return [], 1
                
                client = OpenAI(api_key=api_key, base_url=base_url) if base_url else OpenAI(api_key=api_key)
                
                common_args = {
                    "model": model,
                    # "temperature": 0,
                    "response_format": {"type": "json_object"},
                    # "max_tokens": 1000,
                }
                messages1 = [
                    {"role": "system", "content": msg_system},
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": data_url_in, "detail": "high"}},
                        ],
                    },
                ]
                try:
                    comp = (
                        client.chat.completions.create(messages=messages1, seed=seed, **common_args)
                        if seed is not None
                        else client.chat.completions.create(messages=messages1, **common_args)
                    )
                    if seed is not None:
                        seed_applied = True
                except TypeError:
                    comp = client.chat.completions.create(messages=messages1, **common_args)

                content_loc = comp.choices[0].message.content or "{}"
                result_loc = json.loads(content_loc)

            def _clean_instances(raw: dict[str, Any]) -> list[dict[str, Any]]:
                if "instances" not in raw or not isinstance(raw.get("instances"), list):
                    return []
                cleaned_loc: list[dict[str, Any]] = []
                for inst in raw["instances"]:
                    if not isinstance(inst, dict):
                        continue
                    det = (inst.get("det_prompt") or "").strip()
                    desc = (inst.get("description") or "").strip()
                    bbox = inst.get("bbox")
                    
                    # Validate bbox format
                    if bbox is not None:
                        if not isinstance(bbox, list) or len(bbox) != 4:
                            bbox = None
                        else:
                            try:
                                bbox = [int(x) for x in bbox]
                                # Ensure non-negative dimensions
                                if bbox[2] <= 0 or bbox[3] <= 0:
                                    bbox = None
                            except (ValueError, TypeError):
                                bbox = None
                    
                    if not det:
                        det = desc
                    if not (det or desc):
                        continue
                    if det and not det.rstrip().endswith("."):
                        det = det.rstrip() + " ."
                    
                    cleaned_inst = {
                        "det_prompt": det,
                        "description": desc,
                    }
                    if bbox is not None:
                        cleaned_inst["bbox"] = bbox
                    
                    cleaned_loc.append(cleaned_inst)
                return cleaned_loc

            cleaned_loc = _clean_instances(result_loc)
            attempts_loc = 1
            return cleaned_loc, attempts_loc

        # First, run on the full image
        cleaned_all, attempts_full = _classify_data_url(system_msg, data_url)

        tiles_env = os.getenv("OPENAI_ANALYZER_TILES") or "1"
        try:
            tiles_n = max(1, int(tiles_env))
        except Exception:
            tiles_n = 1

        attempts_total = attempts_full
        used_tiles = 0
        tiles_with_instances = 0
        instances: list[dict[str, Any]] = list(cleaned_all)

        # If configured for tiling and recall is poor, do NxN tiles
        if tiles_n > 1 and not instances:
            try:
                im = Image.open(ipath).convert("RGB")
                W, H = im.size
                tw = max(1, W // tiles_n)
                th = max(1, H // tiles_n)
                for r in range(tiles_n):
                    for c in range(tiles_n):
                        left = c * tw
                        upper = r * th
                        right = min(W, left + tw)
                        lower = min(H, upper + th)
                        tile = im.crop((left, upper, right, lower))
                        # Encode tile to data URL (jpeg)
                        import io
                        buf = io.BytesIO()
                        tile.save(buf, format="JPEG", quality=90)
                        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
                        tile_data_url = f"data:image/jpeg;base64,{b64}"
                        tile_cleaned, tile_attempts = _classify_data_url(system_msg, tile_data_url)
                        attempts_total += tile_attempts
                        used_tiles += 1
                        if tile_cleaned:
                            tiles_with_instances += 1
                            instances.extend(tile_cleaned)
            except Exception:
                pass

        # Deduplicate by (det_prompt, description)
        seen = set()
        deduped: list[dict[str, Any]] = []
        for inst in instances:
            key = (inst.get("det_prompt", ""), inst.get("description", ""))
            if key in seen:
                continue
            seen.add(key)
            deduped.append(inst)

        meta: dict[str, Any] = {
            "provider": provider,
            "model": model,
            "transport": "api",
            "response_format": "json_object",
            "attempts": attempts_total,
            "tiles": tiles_n,
            "tiles_used": used_tiles,
            "tiles_with_instances": tiles_with_instances,
        }
        if base_url and provider == "ollama":
            meta["base_url"] = base_url
        if seed is not None:
            meta["seed"] = seed
            meta["seed_applied"] = bool(seed_applied)

        print(f"[DEBUG] Final return: {len(deduped)} instances", file=sys.stderr, flush=True)
        return {"instances": deduped, "meta": meta}
    except Exception as e:  # pragma: no cover
        print(f"[DEBUG] EXCEPTION: {e}", file=sys.stderr, flush=True)
        return {"instances": [], "error": f"vision model request failed: {e}"}


# Always expose a directly-importable function
def classify_pii(image_path: str) -> Dict[str, Any]:  # noqa: D401
    """Direct function variant (works even if MCP server isn't used)."""
    return _classify_impl(image_path)


if FastMCP is None:
    def main() -> None:  # noqa: D401
        _run_without_mcp()
else:
    app = FastMCP("AnalyzerServer")

    @app.tool()
    def classify_pii(image_path: str) -> Dict[str, Any]:  # type: ignore[override]  # noqa: D401
        """Classify PII instances via OpenAI Vision using prompts/pii_system_prompt.md.

        Output schema (STRICT, per prompt file):
        {
          "instances": [
            {
              "det_prompt": "<dot-terminated detection text>",
              "description": "<rich, PII-aware characterization>",
              "bbox": [x, y, width, height]
            }
          ]
        }
        """
        # Delegate to shared implementation
        return _classify_impl(image_path)

    def main() -> None:  # noqa: D401
        print("[AnalyzerServer] Starting (transport=stdio)…", file=sys.stderr, flush=True)
        try:
            app.run(transport="stdio")
        except BrokenPipeError:
            # Client closed connection during shutdown - this is expected behavior
            pass
        except KeyboardInterrupt:
            # Graceful shutdown on Ctrl+C
            pass
        except BaseException as e:
            # Check if it's an ExceptionGroup containing BrokenPipeError
            if hasattr(e, '__cause__') or hasattr(e, '__context__'):
                cause = getattr(e, '__cause__', None) or getattr(e, '__context__', None)
                if isinstance(cause, BrokenPipeError):
                    pass  # Expected shutdown
                    return
            # Check for ExceptionGroup with BrokenPipeError
            if e.__class__.__name__ in ('ExceptionGroup', 'BaseExceptionGroup'):
                try:
                    exceptions = getattr(e, 'exceptions', [])
                    if any(isinstance(ex, BrokenPipeError) or 
                           'BrokenPipeError' in str(type(ex)) or
                           'Broken pipe' in str(ex) for ex in exceptions):
                        pass  # Expected shutdown
                        return
                except Exception:
                    pass
            print(f"[AnalyzerServer] Unexpected error: {e}", file=sys.stderr, flush=True)
            raise


if __name__ == "__main__":
    main()
