import time
import logging
from typing import Any, Dict
from .config import settings
from .ocr_service import ocr_service, OCRService
from .nvidia_client import nvidia_client, NVIDIAClient
from .prompts import DOCUMENT_INTELLIGENCE_SYSTEM_PROMPT, build_analysis_prompt
from .schemas import DocumentIntelligenceResult

logger = logging.getLogger("document_intelligence.analyzer")


class DocumentAnalyzer:
    """Enterprise document intelligence orchestrator combining OCR and NVIDIA LLM."""

    def __init__(
        self,
        ocr: OCRService = ocr_service,
        llm_client: NVIDIAClient = nvidia_client
    ):
        self.ocr = ocr
        self.llm_client = llm_client

    def analyze(self, image_input: Any) -> DocumentIntelligenceResult:
        """
        End-to-end pipeline:
        Image -> Preprocessing -> OCR Extraction -> Prompt Formatting -> NVIDIA LLM -> Pydantic Schema Validation.
        """
        overall_start = time.time()
        logger.info("Starting document analysis pipeline.")

        # 1. OCR Text Extraction
        try:
            ocr_result = self.ocr.extract_text(image_input)
            ocr_text = ocr_result.get("full_text", "").strip()
            ocr_hints = ocr_result.get("spatial_hints", "")
            ocr_latency = ocr_result.get("latency_seconds", 0.0)
        except Exception as ocr_err:
            logger.error("OCR extract error: %s", ocr_err, exc_info=True)
            return DocumentIntelligenceResult(
                document_type="unknown",
                extracted_text=f"[OCR Processing Error: {str(ocr_err)}]",
                entities=[],
                tables=[],
                components=[],
                relationships=[],
                summary=f"Failed to process image through OCR: {str(ocr_err)}",
                confidence_notes=[f"OCR error details: {str(ocr_err)}", f"Input received type: {type(image_input)}"],
                metadata={"error": str(ocr_err), "input_type": str(type(image_input))}
            )

        if not ocr_text:
            logger.warning("No text detected by OCR engine.")
            return DocumentIntelligenceResult(
                document_type="unknown",
                extracted_text="[No readable text detected in the provided image]",
                entities=[],
                tables=[],
                components=[],
                relationships=[],
                summary="The OCR engine could not identify any readable text or characters in the uploaded image.",
                confidence_notes=[
                    "Low quality image or no text present.",
                    f"Engine attempted: {ocr_result.get('engine', 'unknown')}",
                    f"Image dimensions: {ocr_result.get('image_size', 'unknown')}"
                ],
                metadata={
                    "ocr_engine": ocr_result.get("engine", "unknown"),
                    "ocr_latency_seconds": ocr_latency,
                    "total_latency_seconds": round(time.time() - overall_start, 4)
                }
            )

        # 2. Build Prompt
        user_prompt = build_analysis_prompt(ocr_text, ocr_hints)

        # 3. LLM Structured Extraction
        try:
            llm_response = self.llm_client.generate_document_intelligence(
                system_prompt=DOCUMENT_INTELLIGENCE_SYSTEM_PROMPT,
                user_prompt=user_prompt
            )
            parsed_data = llm_response.get("parsed_data", {})
            llm_latency = llm_response.get("latency_seconds", 0.0)
            model_name = llm_response.get("model", settings.NVIDIA_MODEL)
        except Exception as llm_err:
            logger.warning("LLM call failed: %s. Returning OCR-only response.", llm_err)
            total_latency = time.time() - overall_start
            return DocumentIntelligenceResult(
                document_type="raw_document",
                extracted_text=ocr_text,
                entities=[],
                tables=[],
                components=[],
                relationships=[],
                summary=f"OCR text extraction succeeded ({len(ocr_result.get('lines', []))} lines detected). LLM processing skipped due to: {str(llm_err)}",
                confidence_notes=[
                    f"OCR Engine [{ocr_result.get('engine', 'unknown')}] extracted {len(ocr_result.get('lines', []))} lines in {ocr_latency:.3f}s.",
                    f"NVIDIA LLM Error: {str(llm_err)}. Please set a valid NVIDIA_API_KEY in Space Secrets."
                ],
                metadata={
                    "ocr_engine": ocr_result.get("engine", "unknown"),
                    "ocr_lines_detected": len(ocr_result.get("lines", [])),
                    "ocr_latency_seconds": ocr_latency,
                    "llm_error": str(llm_err),
                    "total_latency_seconds": round(total_latency, 4)
                }
            )

        # 4. Fallback values if LLM skipped some keys
        if "extracted_text" not in parsed_data or not parsed_data["extracted_text"]:
            parsed_data["extracted_text"] = ocr_text

        # 5. Attach execution metadata
        total_latency = time.time() - overall_start
        parsed_data["metadata"] = {
            "ocr_engine": ocr_result.get("engine", "unknown"),
            "ocr_lines_detected": len(ocr_result.get("lines", [])),
            "ocr_latency_seconds": ocr_latency,
            "llm_model": model_name,
            "llm_latency_seconds": llm_latency,
            "total_latency_seconds": round(total_latency, 4),
        }

        # 6. Pydantic validation
        result = DocumentIntelligenceResult(**parsed_data)
        logger.info(
            "Document analysis complete. Type: %s, Entities: %d, Tables: %d, Components: %d in %.3f s",
            result.document_type,
            len(result.entities),
            len(result.tables),
            len(result.components),
            total_latency
        )
        return result


# Global analyzer instance
document_analyzer = DocumentAnalyzer()