"""
Vision analysis service for image understanding (Phase 2B).

Uses the system vision model to analyze extracted images: classify type,
generate descriptions, extract visible text (OCR), and identify concepts
that connect images into the knowledge graph via DEPICTS relationships.
"""

import base64
import json
import logging
import re
import time
from typing import Optional

import requests

from library.metrics import (
    VISION_ANALYSES_TOTAL,
    VISION_ANALYSIS_DURATION,
    VISION_CONCEPTS_EXTRACTED_TOTAL,
)

logger = logging.getLogger(__name__)

# Valid image_type values from the Image model
VALID_IMAGE_TYPES = {
    "cover", "diagram", "chart", "table", "screenshot",
    "illustration", "map", "portrait", "artwork", "still", "photo",
}

# System prompt for structured vision analysis
VISION_SYSTEM_PROMPT = (
    "You are an image analysis assistant. Analyze the image and return a JSON object "
    "with the following fields:\n"
    '- "image_type": one of: cover, diagram, chart, table, screenshot, '
    "illustration, map, portrait, artwork, still, photo\n"
    '- "description": a concise 1-3 sentence description of what the image shows\n'
    '- "ocr_text": any visible text in the image (empty string if none)\n'
    '- "concepts": array of objects with "name" (lowercase) and "type" '
    '(one of: person, place, topic, technique, theme)\n\n'
    "Return ONLY the JSON object, no other text."
)


class VisionAnalyzer:
    """
    Analyzes images using a vision-capable LLM to extract structured metadata.

    For each image, produces:
    - image_type classification
    - natural language description
    - OCR text extraction
    - concept entities for knowledge graph integration

    Requires a system vision model configured in LLM Manager.
    """

    def __init__(self, vision_model, user=None):
        """
        :param vision_model: LLMModel instance for the vision model.
        :param user: Optional Django user for usage tracking.
        """
        self.vision_model = vision_model
        self.api = vision_model.api
        self.user = user
        self.base_url = self.api.base_url.rstrip("/")
        self.model_name = self.vision_model.name
        self.timeout = self.api.timeout_seconds or 120

        logger.info(
            "VisionAnalyzer initialized model=%s api=%s",
            self.model_name,
            self.api.name,
        )

    def analyze_images(
        self,
        image_nodes: list,
        vision_prompt: str = "",
        item=None,
    ) -> int:
        """
        Analyze a list of Image nodes and update them with vision results.

        :param image_nodes: List of Image node instances (with s3_key populated).
        :param vision_prompt: Content-type-aware prompt from the Library config.
        :param item: Optional Item node to connect image concepts via REFERENCES.
        :returns: Number of images successfully analyzed.
        """
        from django.core.files.storage import default_storage

        analyzed_count = 0

        for img_node in image_nodes:
            try:
                # Read image data from S3
                img_data = default_storage.open(img_node.s3_key, "rb").read()
                ext = img_node.s3_key.rsplit(".", 1)[-1] if "." in img_node.s3_key else "png"

                result = self._analyze_single_image(img_data, ext, vision_prompt)

                if result:
                    self._apply_result(img_node, result, item)
                    analyzed_count += 1
                    VISION_ANALYSES_TOTAL.labels(status="success").inc()
                else:
                    img_node.analysis_status = "failed"
                    img_node.save()
                    VISION_ANALYSES_TOTAL.labels(status="failed").inc()

            except Exception as exc:
                logger.warning(
                    "Vision analysis failed for image s3_key=%s: %s",
                    img_node.s3_key,
                    exc,
                )
                img_node.analysis_status = "failed"
                img_node.save()
                VISION_ANALYSES_TOTAL.labels(status="failed").inc()

        if analyzed_count:
            logger.info("Vision analysis completed: %d/%d images", analyzed_count, len(image_nodes))

        return analyzed_count

    def _analyze_single_image(
        self,
        image_data: bytes,
        image_ext: str,
        vision_prompt: str,
    ) -> Optional[dict]:
        """
        Send a single image to the vision model for analysis.

        :param image_data: Raw image bytes.
        :param image_ext: Image format extension (png, jpg, etc.).
        :param vision_prompt: Content-type-aware analysis prompt.
        :returns: Parsed result dict, or None on failure.
        """
        b64 = base64.b64encode(image_data).decode("utf-8")
        mime_type = f"image/{image_ext}" if image_ext != "jpg" else "image/jpeg"

        # Build the user prompt with content-type context
        user_prompt = vision_prompt if vision_prompt else "Analyze this image."

        start_time = time.time()

        try:
            response_text = self._call_vision_model(b64, mime_type, user_prompt)
            elapsed = time.time() - start_time
            VISION_ANALYSIS_DURATION.observe(elapsed)

            result = self._parse_vision_response(response_text)

            logger.debug(
                "Vision analysis completed in %.2fs image_type=%s description_len=%d",
                elapsed,
                result.get("image_type", "unknown") if result else "failed",
                len(result.get("description", "")) if result else 0,
            )
            return result

        except Exception as exc:
            elapsed = time.time() - start_time
            VISION_ANALYSIS_DURATION.observe(elapsed)
            logger.warning("Vision model call failed: %s", exc)
            return None

    def _call_vision_model(
        self,
        b64_image: str,
        mime_type: str,
        user_prompt: str,
    ) -> str:
        """
        Make a chat completion request with an image to the vision model.

        Uses OpenAI-compatible multimodal chat format.

        :param b64_image: Base64-encoded image data.
        :param mime_type: MIME type of the image.
        :param user_prompt: Text prompt to accompany the image.
        :returns: Response text from the model.
        """
        url = f"{self.base_url}/chat/completions"
        headers = {"Content-Type": "application/json"}
        if self.api.api_key:
            headers["Authorization"] = f"Bearer {self.api.api_key}"

        body = {
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": VISION_SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:{mime_type};base64,{b64_image}",
                            },
                        },
                    ],
                },
            ],
            "temperature": 0.1,
            "max_tokens": 800,
        }

        resp = requests.post(url, json=body, headers=headers, timeout=self.timeout)

        if resp.status_code != 200:
            logger.error(
                "Vision model request failed status=%d body=%s",
                resp.status_code,
                resp.text[:500],
            )
        resp.raise_for_status()

        data = resp.json()

        # Parse response — OpenAI-compatible format
        if "choices" in data:
            return data["choices"][0]["message"]["content"]
        if "output" in data:
            # Bedrock Converse format
            return data["output"]["message"]["content"][0]["text"]

        raise ValueError(f"Unexpected vision response format: {list(data.keys())}")

    def _parse_vision_response(self, response_text: str) -> Optional[dict]:
        """
        Parse the vision model's response into structured data.

        :param response_text: Raw response text (expected JSON).
        :returns: Parsed dict with image_type, description, ocr_text, concepts.
        """
        text = response_text.strip()

        # Handle markdown code blocks
        if text.startswith("```"):
            lines = text.split("\n")
            # Remove first line (```json) and last line (```)
            text = "\n".join(lines[1:-1]) if len(lines) > 2 else text

        try:
            result = json.loads(text)
            if isinstance(result, dict):
                return self._validate_result(result)
        except json.JSONDecodeError:
            pass

        # Try to find JSON object in the response
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            try:
                result = json.loads(match.group())
                if isinstance(result, dict):
                    return self._validate_result(result)
            except json.JSONDecodeError:
                pass

        logger.debug("Could not parse vision response: %s", text[:200])
        return None

    def _validate_result(self, result: dict) -> dict:
        """
        Validate and normalize the parsed vision result.

        :param result: Raw parsed dict from the model.
        :returns: Normalized result dict.
        """
        # Normalize image_type
        image_type = result.get("image_type", "").lower().strip()
        if image_type not in VALID_IMAGE_TYPES:
            image_type = "photo"  # Safe default

        # Normalize concepts
        concepts = result.get("concepts", [])
        valid_concepts = []
        for c in concepts:
            if isinstance(c, dict) and "name" in c:
                name = c["name"].strip().lower()
                if name and len(name) >= 2:
                    valid_concepts.append({
                        "name": name,
                        "type": c.get("type", "topic"),
                    })

        return {
            "image_type": image_type,
            "description": str(result.get("description", "")).strip()[:2000],
            "ocr_text": str(result.get("ocr_text", "")).strip()[:5000],
            "concepts": valid_concepts[:20],  # Cap at 20 concepts per image
        }

    def _apply_result(self, img_node, result: dict, item=None):
        """
        Apply vision analysis results to an Image node and create graph relationships.

        :param img_node: Image node to update.
        :param result: Validated result dict from vision analysis.
        :param item: Optional Item node for REFERENCES relationships.
        """
        from library.models import Concept

        # Update Image node fields
        img_node.image_type = result["image_type"]
        img_node.description = result["description"]
        img_node.ocr_text = result["ocr_text"]
        img_node.vision_model_name = self.model_name
        img_node.analysis_status = "completed"
        img_node.save()

        # Create concept relationships
        for concept_data in result["concepts"]:
            name = concept_data["name"]
            concept_type = concept_data["type"]

            concept_node = self._get_or_create_concept(name, concept_type)
            if concept_node:
                # Image DEPICTS Concept
                try:
                    img_node.concepts.connect(concept_node)
                except Exception:
                    pass  # Already connected

                # Also connect Item REFERENCES Concept (if item provided)
                if item:
                    try:
                        item.concepts.connect(concept_node, {"weight": 0.8})
                    except Exception:
                        pass  # Already connected

                VISION_CONCEPTS_EXTRACTED_TOTAL.labels(concept_type=concept_type).inc()

        logger.debug(
            "Applied vision result to image uid=%s type=%s concepts=%d",
            img_node.uid,
            result["image_type"],
            len(result["concepts"]),
        )

        # Log usage
        self._log_usage()

    def _get_or_create_concept(self, name: str, concept_type: str):
        """
        Get or create a Concept node by name (shared with text concept extraction).

        :param name: Concept name (lowercase).
        :param concept_type: Concept type (person, place, topic, etc.).
        :returns: Concept node, or None on failure.
        """
        from library.models import Concept

        try:
            existing = Concept.nodes.filter(name=name)
            if existing:
                return existing[0]

            concept = Concept(name=name, concept_type=concept_type)
            concept.save()
            return concept
        except Exception as exc:
            logger.debug("Failed to get/create concept '%s': %s", name, exc)
            return None

    def _log_usage(self):
        """Log vision analysis usage to LLMUsage model."""
        try:
            from llm_manager.models import LLMUsage

            LLMUsage.objects.create(
                model=self.vision_model,
                user=self.user,
                input_tokens=500,  # Approximate — vision tokens are hard to estimate
                output_tokens=200,
                cached_tokens=0,
                total_cost=0,
                purpose="vision_analysis",
            )
        except Exception as exc:
            logger.warning("Failed to log vision usage: %s", exc)
