"""Semantic tagging service for document extraction pipeline.

Assigns standardized tags (e.g., req:functional, risk:security, topic:authentication)
to extracted records using LLM-based classification with config-driven taxonomy.

Phase 2 Enhancement: Support for ISO-compliant taxonomy + domain profiles.
Phase 3 Enhancement: Tag inheritance, confidence scoring, and JSONB metadata.
"""

import asyncio
import json
from pathlib import Path
from typing import Any, Optional
from uuid import uuid4

import structlog

from src.api.schemas.tags import TagAssignment, TagHistoryEntry, TagSource, TagStatus
from src.core.config import settings
from src.llm import get_chat_service
from src.llm.chat_service import ChatService
from src.extraction_v2.inheritance_engine import InheritanceEngine
from src.extraction_v2.tag_validator import TagValidator
from src.extraction_v2.taxonomy_loader import TaxonomyLoader

logger = structlog.get_logger(__name__)

# TEXTIQ-569: cap on `verified_tag_names` projected onto each record. Mirrors
# the Milvus JSON-path INVERTED index spec; over-cap truncates with a warning.
VERIFIED_TAG_NAMES_MAX_CAPACITY = 64


class PromptBuilder:
    """Builds LLM prompts dynamically from enhanced taxonomy configuration.

    Phase 2 Enhancement: Supports ISO-compliant taxonomy with Hierarchy of Intent.
    """

    def __init__(self, taxonomy_loader: TaxonomyLoader):
        """Initialize prompt builder.

        Args:
            taxonomy_loader: Loaded taxonomy configuration
        """
        self.taxonomy = taxonomy_loader

    def build_system_prompt(self) -> str:
        """Generate system prompt from taxonomy with Hierarchy of Intent (Spec Section 5.2).

        Returns:
            System prompt string with role and taxonomy instructions
        """
        all_tags = self.taxonomy.get_all_tags()
        namespaces = self.taxonomy.taxonomy.get("namespaces", {})
        dynamic_namespaces = self.taxonomy.get_dynamic_namespaces()
        forbidden_namespaces = self.taxonomy.get_forbidden_namespaces()

        # Build tag descriptions grouped by namespace
        tag_groups = {}
        for tag in all_tags:
            ns_name, tag_name = tag.split(":", 1)
            if ns_name not in tag_groups:
                tag_groups[ns_name] = []

            ns_data = namespaces.get(ns_name, {})
            tag_def = next(
                (t for t in ns_data.get("tags", []) if t.get("name") == tag_name),
                None,
            )
            if tag_def:
                desc = tag_def.get("description", "")
                tag_groups[ns_name].append(f"  - {tag} - {desc}")

        # Format pre-defined tags by namespace
        tag_descriptions = []
        for ns_name, tags in tag_groups.items():
            ns_desc = namespaces.get(ns_name, {}).get("description", "")
            tag_descriptions.append(f"**{ns_name}:** {ns_desc}")
            tag_descriptions.extend(tags)
            tag_descriptions.append("")  # Blank line between namespaces

        # Add dynamic namespaces
        tag_descriptions.append("**Dynamic Namespaces (LLM-Generated):**")
        for ns_name in dynamic_namespaces:
            ns_rules = self.taxonomy.namespace_rules["namespaces"][ns_name]
            purpose = ns_rules.get("purpose", "")
            examples = ns_rules.get("examples", [])
            tag_descriptions.append(f"  - {ns_name}:{{keyword}} - {purpose}")
            if examples:
                tag_descriptions.append(f"    Examples: {', '.join(examples[:3])}")

        # Build forbidden list
        forbidden_list = ", ".join(forbidden_namespaces)

        prompt = f"""You are a semantic classification expert trained in ISO 29148, ISO 25010, ISO 42010, and Diátaxis frameworks.

Task: Classify text blocks using the Hierarchy of Intent below.

=== ALLOWED TAGS ===
{chr(10).join(tag_descriptions)}

=== HIERARCHY OF INTENT (CRITICAL) ===
1. GOLDEN RULE (Pre-defined Tags):
   - ALWAYS check if a concept maps to a pre-defined tag first
   - Example: Use "req:quality:security" NOT "topic:security_requirement"

2. SILVER RULE (Standard Dynamic Namespaces):
   - If no pre-defined tag fits, use: {', '.join(dynamic_namespaces)}
   - Match the language of the source content (English lowercase for English docs, native language for others)
   - Examples: "topic:authentication", "tech:postgres" (English); "topic:専門管理加算", "domain:介護" (Japanese)

3. FORBIDDEN (NEVER USE):
   - {forbidden_list}
   - Single words without namespace (e.g., "security" is INVALID)
   - Adjective namespaces (e.g., "important:yes" is INVALID)

=== MULTILINGUAL TAG POLICY ===
- topic:, domain:, and entity: namespaces support native language terms (Japanese, Chinese, etc.)
- For English content: use lowercase with underscores (e.g., "topic:data_privacy")
- For non-English content: use native terms as appropriate (e.g., "topic:専門管理加算", "domain:介護")
- tech: namespace should prefer English standardized names

=== CONFIDENCE SCORING (CRITICAL) ===
You MUST assign a confidence score (0.0-1.0) to each tag using this rubric:
- 1.0 (Certain): Explicitly stated in text (e.g., "The system SHALL validate JWT")
- 0.9 (Very High): Directly implied by unambiguous language
- 0.8 (High): Strongly implied by context or section header
- 0.7 (Moderate-High): Reasonably inferred from surrounding context
- 0.6 (Moderate): Somewhat uncertain but likely correct
- <0.6 (Low): Guessed or very uncertain (DO NOT OUTPUT - discard these)

=== OUTPUT INSTRUCTIONS ===
1. Use parent_header and sheet_name as context
2. Assign multiple tags if applicable (e.g., both functional AND security)
3. Extract 1-3 topic/tech/domain keywords per block
4. Include confidence score for EACH tag
5. ONLY output tags with confidence >= 0.6
6. Return ONLY valid JSON (no markdown, no explanations)
7. If uncertain, use empty array []

Output Format (WITH CONFIDENCE):
[
  {{
    "id": 1,
    "tags": [
      {{"name": "req:functional", "confidence": 1.0}},
      {{"name": "req:quality:security", "confidence": 0.95}},
      {{"name": "topic:authentication", "confidence": 0.85}}
    ]
  }},
  {{
    "id": 2,
    "tags": [
      {{"name": "doc:tutorial", "confidence": 0.9}}
    ]
  }}
]"""
        return prompt

    def build_user_prompt(self, batch: list[dict[str, Any]]) -> str:
        """Build user prompt with batch of records to tag.

        Args:
            batch: List of records with keys: id, text, parent_header, sheet_name

        Returns:
            User prompt string with formatted record batch
        """
        blocks = []
        for record in batch:
            record_id = record.get("id", 0)
            text = record.get("text", "")
            header = record.get("parent_header", "")
            sheet = record.get("sheet_name", "")

            # Truncate long text to avoid prompt bloat
            if len(text) > 500:
                text = text[:497] + "..."

            block = f'Block {record_id} (Header: "{header}", Sheet: "{sheet}"):\n"{text}"'
            blocks.append(block)

        return f"""Input Batch:
{chr(10).join(blocks)}

Return JSON array:"""


class BatchProcessor:
    """Groups records into configurable batches for efficient LLM processing."""

    def __init__(self, batch_size: int | None = None):
        """Initialize batch processor.

        Args:
            batch_size: Number of records per batch (defaults to settings.tag_batch_size)
        """
        self.batch_size = batch_size or settings.tag_batch_size

    def create_batches(
        self, records: list[dict[str, Any]]
    ) -> list[list[dict[str, Any]]]:
        """Split records into batches.

        Args:
            records: List of record dicts with text, parent_header, sheet_name

        Returns:
            List of batches (each batch is a list of records)
        """
        batches = []
        for i in range(0, len(records), self.batch_size):
            batch = records[i : i + self.batch_size]
            batches.append(batch)

        logger.debug(
            "batches_created",
            total_records=len(records),
            num_batches=len(batches),
            batch_size=self.batch_size,
        )

        return batches


class SemanticTagger:
    """Main semantic tagging service.

    Assigns semantic tags to extracted records using LLM-based classification
    with config-driven taxonomy. Integrates with extraction_v2 pipeline.

    Phase 2 Enhancement: Supports domain profiles (fintech, healthcare, automotive).
    """

    def __init__(
        self,
        profile: Optional[str] = None,
        batch_size: int | None = None,
        chat_service: ChatService | None = None,
        auto_approve_threshold: float | None = None,
        # Legacy parameter for backward compatibility
        taxonomy_path: Optional[str | Path] = None,
    ):
        """Initialize semantic tagger.

        Args:
            profile: Domain profile name ("fintech", "healthcare", "automotive", or None for standard)
            batch_size: Records per batch (defaults to settings.tag_batch_size)
            chat_service: Optional AzureChatService instance (for testing/DI)
            auto_approve_threshold: Override the auto-approve confidence
                threshold (tests/DI). Defaults to `settings.tag_auto_approve_threshold`
                (env `TAG_AUTO_APPROVE_THRESHOLD`, default 0.95) — see
                hitl-tag-integration.md §7.
            taxonomy_path: [DEPRECATED] Legacy YAML path (use profile parameter instead)
        """
        # Backward compatibility: if taxonomy_path is provided, use legacy loader
        if taxonomy_path:
            logger.warning(
                "deprecated_parameter",
                parameter="taxonomy_path",
                message="Use 'profile' parameter instead for Phase 2 taxonomy",
            )
            # For legacy YAML files, use old loader
            import yaml
            from pathlib import Path as LegacyPath

            config_path = LegacyPath(taxonomy_path)
            if config_path.exists():
                with open(config_path, encoding="utf-8") as f:
                    legacy_taxonomy = yaml.safe_load(f)
                # Convert to new loader format (minimal compatibility)
                self.taxonomy_loader = TaxonomyLoader()
                self.taxonomy_loader.taxonomy = legacy_taxonomy
            else:
                raise FileNotFoundError(f"Legacy taxonomy not found: {taxonomy_path}")
        else:
            # Use new Phase 2 loader with profile support
            self.taxonomy_loader = TaxonomyLoader(profile=profile)

        self.prompt_builder = PromptBuilder(self.taxonomy_loader)
        self.batch_processor = BatchProcessor(batch_size)
        self.chat_service = chat_service or get_chat_service(temperature=0.0)

        # Phase 3 components
        self.tag_validator = TagValidator(self.taxonomy_loader)
        self.inheritance_engine = InheritanceEngine()
        self.auto_approve_threshold = (
            auto_approve_threshold
            if auto_approve_threshold is not None
            else settings.tag_auto_approve_threshold
        )

        logger.info(
            "semantic_tagger_initialized",
            batch_size=self.batch_processor.batch_size,
            taxonomy_version=self.taxonomy_loader.taxonomy.get("version", "unknown"),
            profile=profile or "standard",
            auto_approve_threshold=self.auto_approve_threshold,
        )

    def tag_records(
        self,
        records: list[dict[str, Any]],
        filename: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> list[dict[str, Any]]:
        """Tag a list of records with semantic tags (Phase 3 Enhanced).

        Args:
            records: List of record dicts with keys:
                - text: Record text content
                - parent_header: Parent section header (for context)
                - sheet_name: Sheet/document name
                - (optional) id: Record identifier
            filename: Document filename for inheritance
            metadata: Document metadata for inheritance

        Returns:
            Same records list with 'tags' field added (list of TagAssignment objects)

        Note:
            - Modifies records in-place by adding 'tags' field
            - On LLM failure, sets tags=[] (empty) and continues
            - Never raises exceptions (fail-safe behavior)
            - Phase 3: Returns TagAssignment objects with confidence, status, history
        """
        if not records:
            return records

        # Add temporary IDs if not present
        for idx, record in enumerate(records):
            if "id" not in record:
                record["id"] = idx

        # Create batches
        batches = self.batch_processor.create_batches(records)
        system_prompt = self.prompt_builder.build_system_prompt()

        # Process each batch
        for batch_num, batch in enumerate(batches, 1):
            try:
                user_prompt = self.prompt_builder.build_user_prompt(batch)

                # Call LLM
                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ]

                response = self.chat_service.invoke(messages)

                # Parse JSON response
                try:
                    tag_results = json.loads(response)
                except json.JSONDecodeError as e:
                    logger.warning(
                        "llm_response_not_json",
                        batch_num=batch_num,
                        error=str(e),
                        response_preview=response[:200],
                    )
                    # Set empty tags for this batch
                    for record in batch:
                        record["tags"] = []
                    continue

                # Apply tags to records (Phase 3: with confidence and validation)
                for tag_result in tag_results:
                    result_id = tag_result.get("id")
                    llm_tags = tag_result.get("tags", [])

                    # Find matching record in batch
                    matching_record = next(
                        (r for r in batch if r.get("id") == result_id), None
                    )
                    if matching_record:
                        # Process LLM tags into TagAssignment objects
                        tag_assignments = self._process_llm_tags(llm_tags)
                        matching_record["tags"] = tag_assignments
                    else:
                        logger.warning(
                            "tag_result_id_mismatch",
                            batch_num=batch_num,
                            result_id=result_id,
                        )

                # Ensure all records in batch have tags field
                for record in batch:
                    if "tags" not in record:
                        record["tags"] = []

                logger.info(
                    "batch_tagged",
                    batch_num=batch_num,
                    batch_size=len(batch),
                    total_batches=len(batches),
                )

            except Exception as e:
                # On any error, set empty tags and continue
                logger.error(
                    "batch_tagging_failed",
                    batch_num=batch_num,
                    batch_size=len(batch),
                    error=str(e),
                    error_type=type(e).__name__,
                )
                for record in batch:
                    record["tags"] = []

        # Apply tag inheritance (document → section → block)
        records = self.inheritance_engine.apply_inheritance(
            records, filename=filename, metadata=metadata
        )

        # TEXTIQ-569: project verified-only names onto each record so the
        # downstream DAG can stamp it onto node.metadata for Milvus's
        # JSON-path INVERTED index (tag-cascade filter).
        self._stamp_verified_tag_names(records)

        logger.info(
            "tagging_complete",
            total_records=len(records),
            total_batches=len(batches),
        )

        return records

    async def _tag_batch_async(
        self,
        semaphore: asyncio.Semaphore,
        batch: list[dict[str, Any]],
        batch_num: int,
        system_prompt: str,
        total_batches: int,
    ) -> None:
        """Tag a single batch with semaphore-controlled concurrency.

        Modifies records in batch in-place by adding 'tags' field.

        Args:
            semaphore: Concurrency limiter
            batch: List of records in this batch
            batch_num: Batch number for logging
            system_prompt: Pre-built system prompt
            total_batches: Total number of batches for logging
        """
        async with semaphore:
            try:
                user_prompt = self.prompt_builder.build_user_prompt(batch)

                messages = [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ]

                response = await self.chat_service.ainvoke(messages)

                try:
                    tag_results = json.loads(response)
                except json.JSONDecodeError as e:
                    logger.warning(
                        "llm_response_not_json",
                        batch_num=batch_num,
                        error=str(e),
                        response_preview=response[:200],
                    )
                    for record in batch:
                        record["tags"] = []
                    return

                for tag_result in tag_results:
                    result_id = tag_result.get("id")
                    llm_tags = tag_result.get("tags", [])

                    matching_record = next(
                        (r for r in batch if r.get("id") == result_id), None
                    )
                    if matching_record:
                        tag_assignments = self._process_llm_tags(llm_tags)
                        matching_record["tags"] = tag_assignments
                    else:
                        logger.warning(
                            "tag_result_id_mismatch",
                            batch_num=batch_num,
                            result_id=result_id,
                        )

                for record in batch:
                    if "tags" not in record:
                        record["tags"] = []

                logger.info(
                    "batch_tagged",
                    batch_num=batch_num,
                    batch_size=len(batch),
                    total_batches=total_batches,
                )

            except Exception as e:
                logger.error(
                    "batch_tagging_failed",
                    batch_num=batch_num,
                    batch_size=len(batch),
                    error=str(e),
                    error_type=type(e).__name__,
                )
                for record in batch:
                    record["tags"] = []

    async def tag_records_async(
        self,
        records: list[dict[str, Any]],
        filename: Optional[str] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> list[dict[str, Any]]:
        """Tag records with concurrent batch processing.

        Same as tag_records but dispatches batches concurrently using asyncio.gather.

        Args:
            records: List of record dicts (see tag_records for format)
            filename: Document filename for inheritance
            metadata: Document metadata for inheritance

        Returns:
            Same records list with 'tags' field added
        """
        if not records:
            return records

        for idx, record in enumerate(records):
            if "id" not in record:
                record["id"] = idx

        batches = self.batch_processor.create_batches(records)
        system_prompt = self.prompt_builder.build_system_prompt()
        concurrency = settings.llm_concurrency
        semaphore = asyncio.Semaphore(concurrency)

        logger.info(
            "async_tagging_starting",
            total_records=len(records),
            total_batches=len(batches),
            concurrency=concurrency,
        )

        tasks = [
            self._tag_batch_async(semaphore, batch, batch_num, system_prompt, len(batches))
            for batch_num, batch in enumerate(batches, 1)
        ]

        await asyncio.gather(*tasks)

        records = self.inheritance_engine.apply_inheritance(
            records, filename=filename, metadata=metadata
        )

        # TEXTIQ-569: project verified-only names onto each record so the
        # downstream DAG can stamp it onto node.metadata for Milvus's
        # JSON-path INVERTED index (tag-cascade filter).
        self._stamp_verified_tag_names(records)

        logger.info(
            "async_tagging_complete",
            total_records=len(records),
            total_batches=len(batches),
            concurrency=concurrency,
        )

        return records

    def _process_llm_tags(self, llm_tags: list[dict[str, Any]]) -> list[TagAssignment]:
        """Process LLM tag output into TagAssignment objects with validation.

        Args:
            llm_tags: List of dicts with 'name' and 'confidence' keys

        Returns:
            List of validated TagAssignment objects
        """
        import time

        current_ts = int(time.time())
        tag_assignments = []

        for llm_tag in llm_tags:
            # Handle both old format (string) and new format (dict with confidence)
            if isinstance(llm_tag, str):
                tag_name = llm_tag
                confidence = 1.0  # Default for legacy format
            elif isinstance(llm_tag, dict):
                tag_name = llm_tag.get("name", "")
                confidence = float(llm_tag.get("confidence", 1.0))
            else:
                logger.warning("invalid_tag_format", tag=llm_tag)
                continue

            if not tag_name:
                continue

            # Validate tag
            is_valid, validation_status, reason = self.tag_validator.validate_tag(
                tag_name
            )

            if not is_valid and validation_status == "rejected":
                logger.warning(
                    "tag_rejected_validation",
                    tag=tag_name,
                    reason=reason,
                )
                continue

            status = self._determine_tag_status(confidence)
            history_action = "confirmed" if status == TagStatus.VERIFIED else "suggested"
            tag_assignment = TagAssignment(
                tag_id=str(uuid4()),
                name=tag_name,
                source=TagSource.AI,
                confidence=confidence,
                status=status,
                history=[
                    TagHistoryEntry(
                        action=history_action,
                        by="ai",
                        ts=current_ts,
                    )
                ],
            )

            tag_assignments.append(tag_assignment)

        return tag_assignments

    def _determine_tag_status(self, confidence: float) -> TagStatus:
        """Resolve initial status for an AI tag.

        Tags with confidence >= `settings.tag_auto_approve_threshold` auto-verify;
        the rest land as `pending`. Operators tune via env
        `TAG_AUTO_APPROVE_THRESHOLD` (default 0.95; set to 1.01 to disable
        auto-approve entirely). See hitl-tag-integration.md §7.
        """
        if confidence >= self.auto_approve_threshold:
            return TagStatus.VERIFIED
        return TagStatus.PENDING

    def _project_verified_tag_names(self, record_tags: Any) -> list[str]:
        """Project names of tags whose status == VERIFIED.

        Output is the flat array stored under the Milvus dynamic-envelope key
        `verified_tag_names` and indexed via JSON-path INVERTED (TEXTIQ-569).
        Accepts list[TagAssignment] (the in-memory shape) or list[dict] (after
        serialization). Caps at VERIFIED_TAG_NAMES_MAX_CAPACITY (64); over-cap
        truncates and warns once.
        """
        if not record_tags:
            return []
        verified: list[str] = []
        for t in record_tags:
            if isinstance(t, dict):
                status = t.get("status")
                name = t.get("name")
            else:
                status = getattr(t, "status", None)
                name = getattr(t, "name", None)
            if hasattr(status, "value"):  # Enum -> str
                status = status.value
            if status != TagStatus.VERIFIED.value or not name:
                continue
            verified.append(name)
        if len(verified) > VERIFIED_TAG_NAMES_MAX_CAPACITY:
            logger.warning(
                "verified_tag_names_truncated",
                dropped=len(verified) - VERIFIED_TAG_NAMES_MAX_CAPACITY,
            )
            verified = verified[:VERIFIED_TAG_NAMES_MAX_CAPACITY]
        return verified

    def _stamp_verified_tag_names(self, records: list[dict[str, Any]]) -> None:
        """Stamp `verified_tag_names` on every record. Idempotent; safe defaults."""
        for record in records:
            record["verified_tag_names"] = self._project_verified_tag_names(
                record.get("tags", [])
            )
