"""Adapter to convert legacy pipeline chunks to new Milvus schema."""

import re
from typing import Any, Dict, List
from uuid import UUID, uuid4
from dataclasses import dataclass

from src.knowledge.embeddings import EmbeddingService


@dataclass
class LegacyChunk:
    """Legacy pipeline chunk structure."""
    file_id: str
    element_type: str
    file_type: str
    text: str
    keywords_text: str
    metadata: Dict[str, Any]
    created_at: str


class LegacyToMilvusAdapter:
    """Converts legacy Excel pipeline chunks to new Milvus schema."""

    def __init__(self):
        self.embedding_service = EmbeddingService()

    async def adapt_chunks(
        self,
        legacy_chunks: List[Dict[str, Any]],
        document_id: UUID,
        filename: str,
    ) -> Dict[str, Any]:
        """Convert legacy chunks to new Milvus entity format.

        Args:
            legacy_chunks: List of chunks from old pipeline
            document_id: Document UUID
            filename: Original filename

        Returns:
            Dict with 'records' and 'embeddings' ready for Milvus
        """
        records = []
        texts_to_embed = []

        for chunk in legacy_chunks:
            # Convert to new format
            record = self._convert_chunk_to_record(chunk, document_id, filename)
            records.append(record)

            # Prepare text for embedding
            text = self._prepare_embedding_text(chunk)
            texts_to_embed.append(text)

        # Generate embeddings in batches (max 100 per batch due to Azure OpenAI limit)
        embeddings = []
        batch_size = 100

        for i in range(0, len(texts_to_embed), batch_size):
            batch = texts_to_embed[i:i + batch_size]
            batch_embeddings = await self.embedding_service.embed_batch(batch)
            embeddings.extend(batch_embeddings)

            if len(texts_to_embed) > batch_size:
                batch_num = (i // batch_size) + 1
                total_batches = (len(texts_to_embed) - 1) // batch_size + 1
                print(f"Generated embeddings for batch {batch_num}/{total_batches} ({len(batch)} chunks)")

        return {
            'records': records,
            'embeddings': embeddings,
            'record_count': len(records)
        }

    def _convert_chunk_to_record(
        self,
        chunk: Dict[str, Any],
        document_id: UUID,
        filename: str,
    ) -> Dict[str, Any]:
        """Convert a single legacy chunk to new record format.

        Maps old schema fields to new ExtractedRecord-like structure.
        """
        element_type = chunk.get('element_type', 'Text')
        metadata = chunk.get('metadata', {})

        # Extract common fields
        sheet_name = metadata.get('sheet_name', 'Sheet1')
        sheet_number = metadata.get('sheet_number', 1)

        # Build content based on element type
        if element_type == 'Text':
            content = self._build_text_content(chunk, metadata)
        elif element_type == 'Image':
            content = self._build_image_content(chunk, metadata)
        elif element_type == 'Shape':
            content = self._build_shape_content(chunk, metadata)
        elif element_type == 'Comment':
            content = self._build_comment_content(chunk, metadata)
        else:
            content = {'text': chunk.get('text', '')}

        # Build _source metadata
        source = {
            'document_id': str(document_id),
            'filename': filename,
            'sheet': sheet_name,
            'table_id': 0,  # Legacy chunks don't have table_id
            'row': self._extract_row_number(metadata, element_type),
            'col_range': self._extract_col_range(metadata, element_type),
        }

        # Create record in ExtractedRecord-like format
        record = {
            'content': content,
            'headers': list(content.keys()),
            '_source': source,
            '_legacy_metadata': {
                'element_type': element_type,
                'keywords_text': chunk.get('keywords_text', ''),
                'created_at': chunk.get('created_at', ''),
                'full_metadata': metadata,
            }
        }

        return record

    def _build_text_content(
        self,
        chunk: Dict[str, Any],
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Build content dict for Text element type."""
        return {
            'text': chunk.get('text', ''),
            'sheet_name': metadata.get('sheet_name', ''),
            'sheet_number': metadata.get('sheet_number', 0),
        }

    def _build_image_content(
        self,
        chunk: Dict[str, Any],
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Build content dict for Image element type."""
        return {
            'description': chunk.get('text', ''),  # AI-generated description
            'image_url': metadata.get('image_url', ''),
            'image_s3_key': metadata.get('image_s3_key', ''),
            'image_format': metadata.get('image_format', ''),
            'width': metadata.get('width', 0),
            'height': metadata.get('height', 0),
            'position_key': metadata.get('position_key', ''),
            'cell_reference': metadata.get('cell_reference', ''),
        }

    def _build_shape_content(
        self,
        chunk: Dict[str, Any],
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Build content dict for Shape element type."""
        connected_target = metadata.get('connected_target', {})
        return {
            'shape_text': chunk.get('text', ''),
            'shape_type': metadata.get('shape_type', ''),
            'target_cell': connected_target.get('cell_ref', ''),
            'target_value': connected_target.get('cell_value', ''),
            'position': str(metadata.get('position', {})),
        }

    def _build_comment_content(
        self,
        chunk: Dict[str, Any],
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Build content dict for Comment element type."""
        return {
            'comment_text': chunk.get('text', ''),
            'comment_id': metadata.get('comment_id', ''),
            'author': metadata.get('author', ''),
            'cell_address': metadata.get('cell_address', ''),
            'cell_value': metadata.get('cell_value', ''),
        }

    def _extract_row_number(
        self,
        metadata: Dict[str, Any],
        element_type: str
    ) -> int:
        """Extract row number from metadata based on element type."""
        if element_type == 'Comment':
            return metadata.get('row', 0)
        elif element_type in ['Image', 'Shape']:
            # Try to extract from position
            position = metadata.get('position', {})
            from_cell = position.get('from_cell', '')
            if from_cell:
                # Extract row number from cell ref (e.g., "A10" -> 10)
                match = re.search(r'\d+', from_cell)
                return int(match.group()) if match else 0
        return 0

    def _extract_col_range(
        self,
        metadata: Dict[str, Any],
        element_type: str
    ) -> str:
        """Extract column range from metadata."""
        if element_type in ['Image', 'Shape']:
            position = metadata.get('position', {})
            from_cell = position.get('from_cell', '')
            to_cell = position.get('to_cell', '')
            if from_cell and to_cell:
                return f"{from_cell}:{to_cell}"
            elif from_cell:
                return from_cell
        elif element_type == 'Comment':
            return metadata.get('cell_address', '')
        return ''

    def _prepare_embedding_text(self, chunk: Dict[str, Any]) -> str:
        """Prepare text for embedding generation.

        Since keyword/entity extraction is removed, simply return the text.
        """
        return chunk.get('text', '')

    def convert_to_milvus_entities(
        self,
        records: List[Dict[str, Any]],
        embeddings: List[List[float]],
        document_id: UUID,
        filename: str,
    ) -> List[Dict[str, Any]]:
        """Convert adapted records to final Milvus entity format.

        This creates the exact structure expected by VectorStore._prepare_milvus_entities.

        Args:
            records: Adapted records (ExtractedRecord-like)
            embeddings: Pre-computed embeddings
            document_id: Document UUID
            filename: Original filename

        Returns:
            List of dicts matching MilvusEntity structure
        """
        entities = []

        for record, embedding in zip(records, embeddings):
            # Extract sheet and row from _source
            source = record.get('_source', {})
            sheet_name = source.get('sheet', 'Sheet1')
            row_number = source.get('row', 0)

            # Check if has symbols (legacy metadata might indicate this)
            legacy_meta = record.get('_legacy_metadata', {})
            has_symbols = legacy_meta.get('element_type') in ['Shape', 'Image', 'Comment']

            # Prepare text_content (truncate for Milvus VARCHAR limit)
            content_text = self._record_to_text(record)
            if len(content_text) > 8000:
                content_text = content_text[:7997] + '...'

            # Build metadata for Milvus - preserve ALL detailed metadata
            # Start with base fields for compatibility
            metadata = {
                'sheet_name': sheet_name,
                'row_number': row_number,
                'has_symbols': has_symbols,
                'element_type': legacy_meta.get('element_type', 'Text'),
            }

            # Preserve ALL additional metadata from DETAILED pipeline
            # The full metadata is stored in _legacy_metadata['full_metadata']
            full_metadata = legacy_meta.get('full_metadata', {})
            if full_metadata:
                for key, value in full_metadata.items():
                    # Skip fields already added (sheet_name, sheet_number) or internal fields
                    if key not in ['sheet_name', 'sheet_number'] and not key.startswith('_'):
                        metadata[key] = value

            entity = {
                'vector': embedding,
                'record_id': str(uuid4()),  # Generate new UUID for each entity
                'document_id': str(document_id),
                'filename': filename,
                'document_type': 'excel',
                'text_content': content_text,
                'metadata': metadata,
            }

            entities.append(entity)

        return entities

    def _record_to_text(self, record: Dict[str, Any]) -> str:
        """Convert record content to text representation for storage."""
        content = record.get('content', {})

        # Format as key: value pairs
        parts = []
        for key, value in content.items():
            if value:
                parts.append(f"{key}: {value}")

        return ' | '.join(parts) if parts else ''