"""Knowledge Base indexer for AiCippy.

Manages S3 storage and Bedrock Knowledge Base indexing.
"""

from __future__ import annotations

import asyncio
import json
import uuid
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

import boto3
from botocore.client import BaseClient
from botocore.config import Config

from aicippy.config import get_settings
from aicippy.utils.logging import get_logger
from aicippy.utils.retry import async_retry

if TYPE_CHECKING:
    from aicippy.knowledge.summarizer import SummarizedContent

logger = get_logger(__name__)


@dataclass
class IndexedDocument:
    """A document indexed in the Knowledge Base."""

    document_id: str
    source: str
    title: str
    s3_key: str
    indexed_at: datetime
    metadata: dict[str, Any]


class KnowledgeIndexer:
    """Indexer for managing Knowledge Base content.

    Handles:
    - S3 storage of documents
    - Knowledge Base data source syncing
    - Document metadata management
    """

    def __init__(self) -> None:
        """Initialize the indexer."""
        self._settings = get_settings()
        self._s3_client: BaseClient | None = None
        self._bedrock_agent_client: BaseClient | None = None

    def _get_s3_client(self) -> BaseClient:
        """Get or create S3 client."""
        if self._s3_client is None:
            config = Config(retries={"max_attempts": 3, "mode": "adaptive"})
            session = boto3.Session(
                region_name=self._settings.aws_region,
                profile_name=self._settings.aws_profile or None,
            )
            self._s3_client = session.client("s3", config=config)
        return self._s3_client

    def _get_bedrock_agent_client(self) -> BaseClient:
        """Get or create Bedrock Agent client."""
        if self._bedrock_agent_client is None:
            config = Config(retries={"max_attempts": 3, "mode": "adaptive"})
            session = boto3.Session(
                region_name=self._settings.aws_region,
                profile_name=self._settings.aws_profile or None,
            )
            self._bedrock_agent_client = session.client("bedrock-agent", config=config)
        return self._bedrock_agent_client

    async def index_documents(
        self,
        documents: list[SummarizedContent],
    ) -> list[IndexedDocument]:
        """Index summarized documents into the Knowledge Base.

        Args:
            documents: List of summarized content to index.

        Returns:
            List of indexed documents.

        """
        logger.info("indexing_started", document_count=len(documents))

        indexed: list[IndexedDocument] = []

        for doc in documents:
            try:
                indexed_doc = await self._index_single_document(doc)
                indexed.append(indexed_doc)
            except Exception as e:
                logger.warning(
                    "document_indexing_failed",
                    title=doc.original_item.title[:50],
                    error=str(e),
                )

        # Trigger Knowledge Base sync
        await self._trigger_kb_sync()

        logger.info("indexing_completed", indexed_count=len(indexed))
        return indexed

    @async_retry(max_attempts=3, min_wait=1.0)
    async def _index_single_document(
        self,
        doc: SummarizedContent,
    ) -> IndexedDocument:
        """Index a single document.

        Args:
            doc: Summarized content to index.

        Returns:
            Indexed document record.

        """
        document_id = str(uuid.uuid4())
        timestamp = datetime.now(UTC)

        # Build document content for S3
        document_content = {
            "title": doc.original_item.title,
            "source": doc.original_item.source,
            "link": doc.original_item.link,
            "published": doc.original_item.published.isoformat() if doc.original_item.published else None,
            "summary": doc.summary,
            "key_points": doc.key_points,
            "categories": doc.categories,
            "original_content": doc.original_item.content[:10000],
            "indexed_at": timestamp.isoformat(),
        }

        # Upload to S3
        s3_key = f"documents/{timestamp.strftime('%Y/%m/%d')}/{document_id}.json"

        s3_client = self._get_s3_client()
        loop = asyncio.get_running_loop()

        await loop.run_in_executor(
            None,
            lambda: s3_client.put_object(
                Bucket=self._settings.s3_knowledge_bucket,
                Key=s3_key,
                Body=json.dumps(document_content, indent=2),
                ContentType="application/json",
                Metadata={
                    "source": doc.original_item.source,
                    "document_id": document_id,
                },
            ),
        )

        # Metadata stored alongside document in S3 (no separate DB needed)
        # Knowledge Base sync handles indexing via Bedrock

        return IndexedDocument(
            document_id=document_id,
            source=doc.original_item.source,
            title=doc.original_item.title,
            s3_key=s3_key,
            indexed_at=timestamp,
            metadata={
                "categories": doc.categories,
                "key_points": doc.key_points,
            },
        )

    async def _trigger_kb_sync(self) -> None:
        """Trigger Knowledge Base data source sync."""
        client = self._get_bedrock_agent_client()
        if client is None:
            logger.warning("kb_sync_skipped_no_client")
            return
        try:
            kb_id = self._settings.bedrock_knowledge_base_id
            ds_id = self._settings.bedrock_data_source_id
            if not kb_id or not ds_id:
                logger.warning("kb_sync_skipped_no_config", kb_id=kb_id, ds_id=ds_id)
                return
            await asyncio.get_running_loop().run_in_executor(
                None,
                lambda: client.start_ingestion_job(knowledgeBaseId=kb_id, dataSourceId=ds_id),
            )
            logger.info("kb_sync_started", kb_id=kb_id)
        except Exception as e:
            logger.error("kb_sync_failed", error=str(e))

    async def search(
        self,
        query: str,
        max_results: int = 10,
    ) -> list[dict[str, Any]]:
        """Search the Knowledge Base.

        Args:
            query: Search query.
            max_results: Maximum results to return.

        Returns:
            List of matching documents.

        """
        client = self._get_bedrock_agent_client()
        if client is None:
            logger.warning("kb_search_skipped_no_client")
            return []
        try:
            kb_id = self._settings.bedrock_knowledge_base_id
            if not kb_id:
                return []
            response = await asyncio.get_running_loop().run_in_executor(
                None,
                lambda: client.retrieve(
                    knowledgeBaseId=kb_id,
                    retrievalQuery={"text": query},
                    retrievalConfiguration={
                        "vectorSearchConfiguration": {"numberOfResults": max_results},
                    },
                ),
            )
            results = []
            for item in response.get("retrievalResults", []):
                results.append({
                    "content": item.get("content", {}).get("text", ""),
                    "score": item.get("score", 0.0),
                    "location": item.get("location", {}),
                })
            return results
        except Exception as e:
            logger.error("kb_search_failed", error=str(e), query=query[:50])
            return []

    async def delete_document(self, document_id: str) -> bool:
        """Delete a document from the Knowledge Base.

        Args:
            document_id: Document ID to delete.

        Returns:
            True if deleted, False on failure.

        """
        s3 = self._get_s3_client()
        if s3 is None:
            logger.warning("delete_document_skipped_no_client")
            return False
        try:
            bucket = self._settings.s3_knowledge_bucket
            key = f"documents/{document_id}"
            await asyncio.get_running_loop().run_in_executor(
                None, lambda: s3.delete_object(Bucket=bucket, Key=key),
            )
            logger.info("document_deleted", document_id=document_id)
            await self._trigger_kb_sync()
            return True
        except Exception as e:
            logger.error("document_delete_failed", error=str(e))
            return False
