import json
import time
from collections import defaultdict
from typing import Any, List

from google.api_core import exceptions as gcloud_exceptions
from google.api_core.client_options import ClientOptions
from google.cloud import documentai  # type: ignore
from google.protobuf.json_format import MessageToDict
from pydantic import ValidationError

from .exceptions import OcrError, OcrFailureCode
from .models import NormalizedValue, OCRResult, RawEntity, RawOcrDocument
from .ocr_processors import OCRProcessorName, OCR_PROCESSORS, OCRProcessor
from .parse_ocr.parsers import OCR_PARSERS
from .parse_ocr.wtp_left_page_parser import WTPLeftPageOCR
from .parse_ocr.wtp_right_page_parser import WTPRightPageOCR
from .parse_ocr.wtp_coagulation_parser import WTPCoagulationLogOCR

# Document AI errors that mean "the request/credentials/processor config were
# bad", not "the file was bad" -- everything else in GoogleAPICallError (timeouts,
# quota, service unavailable, transient internal errors) is treated as a
# provider-side failure below.
_DOCUMENT_AI_CONFIGURATION_ERRORS = (
    gcloud_exceptions.PermissionDenied,
    gcloud_exceptions.Unauthenticated,
    gcloud_exceptions.NotFound,
)


def _text_from_anchor(
    text_anchor: documentai.Document.TextAnchor, full_text: str
) -> str:
    """Build text for an anchor from the full document text."""
    if not text_anchor or not text_anchor.text_segments:
        return ""

    parts: List[str] = []
    for segment in text_anchor.text_segments:
        start_index = int(segment.start_index) if segment.start_index else 0
        end_index = int(segment.end_index)
        parts.append(full_text[start_index:end_index])
    return "".join(parts)


def _normalized_value_from_proto(message: Any) -> NormalizedValue | None:
    """Convert a Document AI normalized value proto into a model."""
    if not message:
        return None
    if hasattr(message, "_pb"):
        data = MessageToDict(message._pb, preserving_proto_field_name=True)
        return NormalizedValue.model_validate(data)
    return None


def _entity_from_document(
    entity: documentai.Document.Entity, full_text: str
) -> RawEntity:
    """Convert a Document AI entity into a RawEntity."""
    mention_text = entity.mention_text or _text_from_anchor(entity.text_anchor, full_text)

    return RawEntity(
        type=entity.type_,
        mention_text=mention_text,
        normalized_value=_normalized_value_from_proto(entity.normalized_value),
        confidence=entity.confidence,
        id=entity.id or None,
        properties=[_entity_from_document(prop, full_text) for prop in entity.properties],
    )


def _get_document_ocr(
    ocr_processor: OCRProcessor,
    file_path: str,
    mime_type: str,
) -> RawOcrDocument:
    # You must set the `api_endpoint` if you use a location other than "us".
    opts = ClientOptions(api_endpoint=f"{ocr_processor.location}-documentai.googleapis.com")

    client = documentai.DocumentProcessorServiceClient(client_options=opts)
    name = client.processor_path(ocr_processor.project_id, ocr_processor.location, ocr_processor.processor_id)

    # Read the file into memory
    with open(file_path, "rb") as image:
        image_content = image.read()

    # Load binary datao
    raw_document = documentai.RawDocument(content=image_content, mime_type=mime_type)
    process_options = documentai.ProcessOptions(
        # Process only specific pages
        individual_page_selector=documentai.ProcessOptions.IndividualPageSelector(
            pages=[1]
        )
    )

    # Configure the process request
    request = documentai.ProcessRequest(
        name=name,
        raw_document=raw_document,
        process_options=process_options,
    )

    try:
        result = client.process_document(request=request)
    except _DOCUMENT_AI_CONFIGURATION_ERRORS as exc:
        raise OcrError(OcrFailureCode.CONFIGURATION_ERROR, f"Document AI rejected the request: {exc}") from exc
    except gcloud_exceptions.InvalidArgument as exc:
        raise OcrError(OcrFailureCode.UNSUPPORTED_FILE, f"Document AI rejected the file: {exc}") from exc
    except gcloud_exceptions.GoogleAPICallError as exc:
        raise OcrError(OcrFailureCode.PROVIDER_ERROR, f"Document AI call failed: {exc}") from exc

    # For a full list of `Document` object attributes, reference this page:
    # https://cloud.google.com/document-ai/docs/reference/rest/v1/Document
    document = result.document

    entity_items = [
        _entity_from_document(entity, document.text) for entity in document.entities
    ]
    entities_by_type: dict[str, list[RawEntity]] = defaultdict(list)
    for entity in entity_items:
        entities_by_type[entity.type].append(entity)

    return RawOcrDocument(
        project_id=ocr_processor.project_id,
        location=ocr_processor.location,
        processor_id=ocr_processor.processor_id,
        source_file=file_path,
        ocr_text=document.text,
        entities=entity_items,
        entities_by_type=dict(entities_by_type),
    )


def _get_structured_ocr(
    file_path: str,
    mime_type: str,
    processor_name: OCRProcessorName,
    raw_document_ocr_path: str | None = None
) -> OCRResult:
    start = time.perf_counter()

    ocr_processor = OCR_PROCESSORS[processor_name]
    raw_ocr_document = _get_document_ocr(ocr_processor, file_path, mime_type)

    if raw_document_ocr_path is not None:
        with open(raw_document_ocr_path, "w") as f:
            json.dump(raw_ocr_document.model_dump(), f)
        
    parser = OCR_PARSERS[processor_name]()
    try:
        result = parser.parse(raw_ocr_document)
    except ValidationError as exc:
        raise OcrError(
            OcrFailureCode.EXTRACTION_VALIDATION_FAILED,
            f"Extracted data failed validation: {exc}",
        ) from exc

    return OCRResult(
        result=result,
        time_elapsed=time.perf_counter() - start,
    )


def get_wtp_left_page_ocr(
    file_path: str,
    mime_type: str,
    raw_document_ocr_path: str | None = None
) -> OCRResult[WTPLeftPageOCR]:
    return _get_structured_ocr(
        file_path,
        mime_type,
        OCRProcessorName.WTP_LEFT_PAGE,
        raw_document_ocr_path
    )

def get_wtp_right_page_ocr(
    file_path: str,
    mime_type: str,
    raw_document_ocr_path: str | None = None
) -> OCRResult[WTPRightPageOCR]:
    return _get_structured_ocr(
        file_path,
        mime_type,
        OCRProcessorName.WTP_RIGHT_PAGE,
        raw_document_ocr_path
    )

def get_wtp_coagulation_log_ocr(
    file_path: str,
    mime_type: str,
    raw_document_ocr_path: str | None = None
) -> OCRResult[WTPCoagulationLogOCR]:
    return _get_structured_ocr(
        file_path,
        mime_type,
        OCRProcessorName.WTP_COAGULATION,
        raw_document_ocr_path
    )

