#  Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
#  This file is part of Archive Agent. See LICENSE for details.

import json
from enum import Enum
from typing import cast, Dict, List, Optional

import typer
from openai import BadRequestError
from qdrant_client.http.models import ScoredPoint

from archive_agent.ai.AiResult import AiResult
from archive_agent.ai.chunk.AiChunk import AiChunk, ChunkSchema
from archive_agent.ai.query.AiQuery import AiQuery, QuerySchema
from archive_agent.ai.rerank.AiRerank import AiRerank, RerankSchema
from archive_agent.ai.vision.AiVisionEntity import AiVisionEntity
from archive_agent.ai.vision.AiVisionOCR import AiVisionOCR
from archive_agent.ai.vision.AiVisionSchema import VisionSchema, VisionSchemaMultiline
from archive_agent.ai_provider.AiProvider import AiProvider
from archive_agent.ai_provider.AiProviderError import AiProviderMaxTokensError

from archive_agent.core.CliManager import CliManager
from archive_agent.util.RetryManager import RetryManager
from archive_agent.util.text_util import prepend_line_numbers


class AiVisionRequest(Enum):
    ENTITY = 'entity'
    OCR = 'ocr'
    OCR_STANDALONE = 'ocr_standalone'


class AiManager(RetryManager):
    """
    AI manager.
    """

    AI_RETRY_KWARGS = {
        'predelay': 0.0,
        'delay_min': 0.0,
        'delay_max': 60.0,
        'backoff_exponent': 2.0,
        'retries': 10,
    }

    SCHEMA_RETRY_ATTEMPTS = 10
    EMBED_TRUNCATION_ATTEMPTS = 10

    def __init__(
            self,
            cli: CliManager,
            chunk_lines_block: int,
            chunk_words_target: int,
            ai_provider: AiProvider,
    ):
        """
        Initialize AI manager.
        :param cli: CLI manager.
        :param chunk_lines_block: Number of lines per block for chunking.
        :param chunk_words_target: Target number of words per chunk.
        :param ai_provider: AI provider.
        """
        self.cli = cli

        self.chunk_lines_block = chunk_lines_block
        self.chunk_words_target = chunk_words_target

        self.ai_provider = ai_provider

        self.ai_usage_stats = {
            "chunk": 0,
            "embed": 0,
            "rerank": 0,
            "query": 0,
            "vision": 0
        }

        # NOTE: This switches between `AiVisionEntity` and `AiVisionOCR` modules
        self.requested: Optional[AiVisionRequest] = None

        RetryManager.__init__(self, **AiManager.AI_RETRY_KWARGS)

        if not self.ai_provider.supports_vision:
            self.cli.logger.warning(f"Image vision is DISABLED in your current configuration")

    def usage(self):
        """
        Show AI token usage.
        """
        if any(self.ai_usage_stats.values()):
            self.cli.logger.info(
                f"Used AI API token(s): "
                f"({self.ai_usage_stats['chunk']}) chunking, "
                f"({self.ai_usage_stats['embed']}) embedding, "
                f"({self.ai_usage_stats['rerank']}) reranking, "
                f"({self.ai_usage_stats['query']}) query, "
                f"({self.ai_usage_stats['vision']}) vision"
            )
        else:
            self.cli.logger.info(f"No AI API tokens used")

    def chunk(self, sentences: List[str]) -> ChunkSchema:
        """
        Get chunks of sentences.
        :param sentences: Sentences.
        :return: ChunkSchema.
        """
        line_numbered_text = "\n".join(prepend_line_numbers(sentences))
        prompt = AiChunk.get_prompt_chunk(line_numbered_text=line_numbered_text, chunk_words_target=self.chunk_words_target)
        callback = lambda: self.ai_provider.chunk_callback(prompt=prompt)

        for _ in range(AiManager.SCHEMA_RETRY_ATTEMPTS):
            try:
                result: AiResult = self.cli.format_ai_chunk(callback=lambda: self.retry(callback), line_numbered_text=line_numbered_text)
                self.ai_usage_stats['chunk'] += result.total_tokens

                if result.parsed_schema is None:
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError("No parsed schema returned")

                chunk_result = cast(ChunkSchema, result.parsed_schema)
                chunk_start_lines = chunk_result.get_chunk_start_lines()
                chunk_headers = chunk_result.get_chunk_headers()

                if len(chunk_start_lines) != len(chunk_headers):
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError(
                        f"Mismatch: "
                        f"chunk_start_lines[{len(chunk_start_lines)}] != headers[{len(chunk_headers)}]"
                    )

                if len(chunk_start_lines) == 0:
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError(f"Missing chunk start lines: {chunk_start_lines}")

                if chunk_start_lines[0] != 1:
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError(f"First chunk must start at line 1: {chunk_start_lines}")

                if not 0 < any(chunk_start_lines) <= len(sentences):
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError(f"Invalid line numbers: {chunk_start_lines}")

                return chunk_result

            except typer.Exit:
                raise  # Network retries exhausted — don't swallow process exit
            except AiProviderMaxTokensError:
                raise  # Non-retryable — retrying same input will always truncate
            except Exception as e:
                self.cli.logger.exception(f"Chunking error: {e}")
                continue  # Retry

        self.ai_provider.invalidate_last_cached()
        raise RuntimeError("Failed to recover from chunking errors")

    def embed(self, text: str) -> List[float]:
        """
        Embed text.
        :param text: Text.
        :return: Embedding vector.
        """
        original_text = text
        original_word_count = len(text.split())

        for truncation_attempt in range(AiManager.EMBED_TRUNCATION_ATTEMPTS + 1):
            callback = lambda: self.ai_provider.embed_callback(text)

            try:
                result: AiResult = self.cli.format_ai_embed(callback=lambda: self.retry(callback), text=text)
                self.ai_usage_stats['embed'] += result.total_tokens
                assert result.embedding is not None
                return result.embedding
            except BadRequestError as e:
                error_body = getattr(e, 'body', {})
                error_message = error_body.get('message', '') if isinstance(error_body, dict) else str(e)

                # Check for token limit error (rare case: chunk exceeds embedding model token limit)
                if "maximum context length" in error_message and truncation_attempt < AiManager.EMBED_TRUNCATION_ATTEMPTS:
                    # Progressive truncation: 10% per attempt (last resort for oversized chunks)
                    truncation_percent = (truncation_attempt + 1) * 10
                    keep_percent = 100 - truncation_percent

                    # Truncate by words
                    words = original_text.split()
                    truncated_word_count = int(len(words) * keep_percent / 100)
                    text = " ".join(words[:truncated_word_count])

                    self.cli.logger.warning(
                        f"⚠️ Chunk exceeded embedding token limit (truncation attempt {truncation_attempt + 1}). "
                        f"Truncating {truncation_percent}% ({original_word_count} → {truncated_word_count} words) and retrying... "
                        f"Please check your config to avoid this error."
                    )
                    continue  # Retry with truncated text

                # Re-raise if not a token limit error or max attempts reached
                raise

        raise RuntimeError(f"Failed to embed after {AiManager.EMBED_TRUNCATION_ATTEMPTS} truncation attempts")

    def rerank(self, question: str, indexed_chunks: Dict[int, str]) -> RerankSchema:
        """
        Get reranked chunks based on relevance to question.
        :param question: Question.
        :param indexed_chunks: Indexed chunks.
        :return: RerankSchema.
        """
        indexed_chunks_json_text = json.dumps(indexed_chunks, ensure_ascii=False, indent=2)
        prompt = AiRerank.get_prompt_rerank(question=question, indexed_chunks_json_text=indexed_chunks_json_text)
        callback = lambda: self.ai_provider.rerank_callback(prompt=prompt)

        for _ in range(AiManager.SCHEMA_RETRY_ATTEMPTS):
            try:
                result: AiResult = self.cli.format_ai_rerank(callback=lambda: self.retry(callback), indexed_chunks=indexed_chunks)
                self.ai_usage_stats['rerank'] += result.total_tokens

                if result.parsed_schema is None:
                    self.ai_provider.invalidate_last_cached()
                    raise RuntimeError("No parsed schema returned")

                rerank_result = cast(RerankSchema, result.parsed_schema)

                ai_is_stupid = rerank_result.reranked_indices == [0]  # Let's allow some slack from weaker or overloaded LLMs here...
                if rerank_result.is_rejected or ai_is_stupid:
                    self.ai_provider.invalidate_last_cached()
                    self.cli.logger.critical(f"⚠️ Reranking context rejected: \"{rerank_result.rejection_reason}\"")
                    return rerank_result

                reranked = rerank_result.reranked_indices
                expected = list(indexed_chunks.keys())

                # First, validate strictly
                is_valid, missing, extra, oob = AiRerank.validate_permutation(expected, reranked)
                if not is_valid:
                    self.ai_provider.invalidate_last_cached()

                    # Auto-repair path for production solidity
                    repaired = AiRerank.repair_permutation(expected, reranked)

                    # Validate again to be absolutely sure
                    is_valid2, missing2, extra2, oob2 = AiRerank.validate_permutation(expected, repaired)
                    if not is_valid2:
                        raise RuntimeError(
                            "Reranked indices could not be repaired into a valid permutation:\n"
                            f"Original: {expected}\n"
                            f"Reranked: {reranked}\n"
                            f"Missing: {missing} | Extra: {extra} | Out-of-range: {oob}\n"
                            f"Repair missing: {missing2} | Repair extra: {extra2} | Repair out-of-range: {oob2}"
                        )

                    # Log the discrepancy and proceed with repaired indices
                    self.cli.logger.warning(
                        "⚠️ Auto-repaired invalid rerank permutation.\n"
                        f"Original: {expected}\n"
                        f"Received: {reranked}\n"
                        f"Repaired: {repaired}\n"
                        f"Missing: {missing} | Extra: {extra} | Out-of-range: {oob}"
                    )
                    # Mutate the schema in-place so downstream users see the corrected order.
                    rerank_result.reranked_indices = repaired

                return rerank_result

            except typer.Exit:
                raise  # Network retries exhausted — don't swallow process exit
            except AiProviderMaxTokensError:
                raise  # Non-retryable — retrying same input will always truncate
            except Exception as e:
                self.cli.logger.exception(f"Reranking error: {e}")
                continue  # Retry

        self.ai_provider.invalidate_last_cached()
        raise RuntimeError("Failed to recover from reranking errors")

    def query(self, question: str, points: List[ScoredPoint]) -> QuerySchema:
        """
        Get answer to question using RAG.
        :param question: Question.
        :param points: Points.
        :return: QuerySchema.
        """
        context = AiQuery.get_context_from_points(points)
        prompt = AiQuery.get_prompt_query(question=question, context=context)
        callback = lambda: self.ai_provider.query_callback(prompt=prompt)

        result: AiResult = self.cli.format_ai_query(callback=lambda: self.retry(callback), prompt=prompt)
        self.ai_usage_stats['query'] += result.total_tokens
        assert result.parsed_schema is not None
        query_result = cast(QuerySchema, result.parsed_schema)
        query_result = AiQuery.format_query_references(logger=self.cli.logger, query_result=query_result, points=points)

        if query_result.is_rejected:
            self.ai_provider.invalidate_last_cached()

        return query_result

    def vision(self, image_base64: str) -> VisionSchema:
        """
        Convert image to text.
        :param image_base64: Image as UTF-8 encoded Base64 string.
        :return: VisionSchema.
        """
        is_standalone = self.requested == AiVisionRequest.OCR_STANDALONE

        if self.requested == AiVisionRequest.ENTITY:
            prompt = AiVisionEntity.get_prompt_vision()
        elif self.requested == AiVisionRequest.OCR:
            prompt = AiVisionOCR.get_prompt_vision()
        elif self.requested == AiVisionRequest.OCR_STANDALONE:
            prompt = AiVisionOCR.get_prompt_vision_standalone()
        else:
            self.cli.logger.critical("⚠️ BUG DETECTED: Unrequested call to `AiManager.vision()` — falling back to OCR")
            prompt = AiVisionOCR.get_prompt_vision()

        self.requested = None

        callback = lambda: self.ai_provider.vision_callback(prompt=prompt, image_base64=image_base64)

        result: AiResult = self.cli.format_ai_vision(callback=lambda: self.retry(callback))
        self.ai_usage_stats['vision'] += result.total_tokens
        assert result.parsed_schema is not None
        vision_result = cast(VisionSchema, result.parsed_schema)

        if is_standalone and not vision_result.is_rejected:
            # Standalone OCR ONLY: re-parse the raw JSON with the multiline schema to
            # preserve structural line breaks. The base VisionSchema (used everywhere else)
            # collapses them to guarantee single-line output for the RAG pipeline.
            vision_result = VisionSchemaMultiline.model_validate_json(
                self.ai_provider._sanitize_json(result.output_text)
            )

        if vision_result.is_rejected:
            self.ai_provider.invalidate_last_cached()

        return vision_result

    def request_entity(self):
        self.requested = AiVisionRequest.ENTITY

    def request_ocr(self):
        self.requested = AiVisionRequest.OCR

    def request_ocr_standalone(self):
        self.requested = AiVisionRequest.OCR_STANDALONE
