"""
Entity Resolution Service

Intelligently selects the best agent or team from an organization's catalog
based on task requirements using LLM analysis.

Similar to queue_resolution_service but for selecting execution entities.
"""

from typing import Dict, Any, List, Optional
import structlog
import re
import json
import asyncio
import time
from functools import lru_cache

from sqlalchemy.orm import Session
from control_plane_api.app.models.agent import Agent
from control_plane_api.app.models.team import Team
from control_plane_api.app.services.agno_service import agno_service

logger = structlog.get_logger()


class EntityResolutionService:
    """Service for intelligent agent/team selection using LLM analysis"""

    def __init__(self):
        # Cache for entity lists (short TTL to stay fresh)
        self._entity_cache: Dict[str, tuple[List[Dict], float]] = {}
        self._entity_cache_ttl = 60  # 60 seconds

        # Performance metrics
        self._metrics = {
            "llm_calls": 0,
            "llm_failures": 0,
            "pattern_only_selections": 0,
            "cache_hits": 0,
        }

    async def resolve_entity(
        self,
        db: Session,
        prompt: str,
        organization_id: str,
        entity_type: str = "agent",  # "agent" or "team"
    ) -> Dict[str, Any]:
        """
        Analyze task prompt and recommend optimal agent or team.

        Args:
            db: Database session
            prompt: Task prompt to analyze
            organization_id: Organization ID
            entity_type: Type of entity to select ("agent" or "team")

        Returns:
            {
                "entity_id": str,
                "entity_name": str,
                "entity_type": str,
                "reasoning": str,
                "confidence": "high" | "medium" | "low",
                "alternatives": List[Dict]
            }
        """
        start_time = time.time()

        logger.info(
            "Resolving entity",
            organization_id=organization_id,
            entity_type=entity_type,
            prompt_length=len(prompt)
        )

        try:
            # STEP 1: Fetch available entities (with caching)
            available_entities = await self._get_available_entities_cached(
                db=db,
                organization_id=organization_id,
                entity_type=entity_type
            )

            if not available_entities:
                raise ValueError(f"No {entity_type}s found for organization {organization_id}")

            # STEP 2: Analyze task requirements
            task_analysis = self._analyze_task_patterns(prompt)

            logger.info(
                "Task analysis complete",
                task_category=task_analysis.get("task_category"),
                keywords=task_analysis.get("keywords"),
                complexity=task_analysis.get("complexity")
            )

            # STEP 3: Score entities based on analysis
            scored_entities = self._score_entities(task_analysis, available_entities)

            # STEP 4: Use LLM for final recommendation if scores are close
            if len(scored_entities) >= 2 and abs(scored_entities[0]["score"] - scored_entities[1]["score"]) <= 15:
                logger.info("Scores are close, using LLM for final recommendation")

                try:
                    llm_recommendation = await self._llm_recommend_entity_with_retry(
                        prompt=prompt,
                        task_analysis=task_analysis,
                        top_entities=scored_entities[:3],  # Top 3 candidates
                        entity_type=entity_type
                    )

                    # Find the recommended entity in scored_entities
                    recommended = next(
                        (e for e in scored_entities if e["entity_id"] == llm_recommendation["entity_id"]),
                        scored_entities[0]  # Fallback to highest scored
                    )

                    # Update reasoning with LLM insight
                    recommended["reasoning"] = f"{recommended['reasoning']}; LLM: {llm_recommendation.get('reasoning', '')}"

                except Exception as e:
                    logger.warning("LLM recommendation failed, using pattern-based", error=str(e))
                    recommended = scored_entities[0]
                    self._metrics["pattern_only_selections"] += 1
            else:
                # Clear winner, use pattern-based recommendation
                recommended = scored_entities[0]
                self._metrics["pattern_only_selections"] += 1

            elapsed = time.time() - start_time
            logger.info(
                "Entity resolution complete",
                entity_id=recommended["entity_id"],
                entity_name=recommended["entity_name"],
                entity_type=entity_type,
                confidence=recommended["confidence"],
                score=recommended["score"],
                elapsed_ms=int(elapsed * 1000)
            )

            return {
                "entity_id": recommended["entity_id"],
                "entity_name": recommended["entity_name"],
                "entity_type": entity_type,
                "reasoning": recommended["reasoning"],
                "confidence": recommended["confidence"],
                "alternatives": scored_entities[1:3]  # Next 2 alternatives
            }

        except ValueError as e:
            # Re-raise validation errors
            logger.error("Entity resolution validation error", error=str(e))
            raise
        except Exception as e:
            logger.error("Entity resolution failed", error=str(e), exc_info=True)
            raise RuntimeError(f"Failed to resolve {entity_type}: {str(e)}")

    async def _get_available_entities_cached(
        self,
        db: Session,
        organization_id: str,
        entity_type: str
    ) -> List[Dict[str, Any]]:
        """
        Fetch available agents or teams from database with caching.

        Returns:
            List of entity dicts with metadata
        """
        cache_key = f"{organization_id}:{entity_type}"

        # Check cache
        if cache_key in self._entity_cache:
            cached_entities, cached_time = self._entity_cache[cache_key]
            if time.time() - cached_time < self._entity_cache_ttl:
                self._metrics["cache_hits"] += 1
                logger.debug("Entity cache hit", cache_key=cache_key)
                return cached_entities

        # Cache miss - fetch from database
        entities = await self._get_available_entities(db, organization_id, entity_type)

        # Update cache
        self._entity_cache[cache_key] = (entities, time.time())

        # Clean old cache entries
        self._cleanup_cache()

        return entities

    async def _get_available_entities(
        self,
        db: Session,
        organization_id: str,
        entity_type: str
    ) -> List[Dict[str, Any]]:
        """
        Fetch available agents or teams from database.

        Returns:
            List of entity dicts with metadata
        """
        try:
            if entity_type == "agent":
                Model = Agent
            elif entity_type == "team":
                Model = Team
            else:
                raise ValueError(f"Invalid entity_type: {entity_type}")

            entities = db.query(Model).filter(
                Model.organization_id == organization_id,
            ).all()

            result = []
            for entity in entities:
                try:
                    entity_dict = {
                        "id": str(entity.id),
                        "name": entity.name,
                        "description": entity.description or "",
                        "system_prompt": getattr(entity, "system_prompt", ""),
                        "tags": getattr(entity, "tags", []) or [],
                        "created_at": entity.created_at.isoformat() if entity.created_at else None,
                    }

                    # For teams, add subagent count
                    if entity_type == "team":
                        # Assuming team has a configuration field with subagents
                        config = getattr(entity, "configuration", {}) or {}
                        runtime_config = config.get("runtime_config", {})
                        agents_config = runtime_config.get("agents", {})
                        entity_dict["subagent_count"] = len(agents_config)

                    result.append(entity_dict)

                except Exception as e:
                    logger.warning(
                        "Failed to process entity",
                        entity_id=str(entity.id),
                        error=str(e)
                    )
                    continue

            logger.info(
                "Fetched available entities",
                organization_id=organization_id,
                entity_type=entity_type,
                count=len(result)
            )

            return result

        except Exception as e:
            logger.error("Failed to fetch entities from database", error=str(e))
            raise

    @lru_cache(maxsize=256)
    def _analyze_task_patterns(self, prompt: str) -> Dict[str, Any]:
        """
        Analyze task prompt using pattern matching to extract requirements.
        Cached for performance.

        Returns:
            {
                "task_category": str,
                "complexity": str,
                "keywords": List[str],
                "prompt_length": int
            }
        """
        prompt_lower = prompt.lower()

        # Categorize task
        category_patterns = {
            "deployment": r"\b(deploy|release|rollout|publish|ship)\b",
            "infrastructure": r"\b(provision|infrastructure|setup|terraform|cloudformation)\b",
            "development": r"\b(code|develop|implement|build|create|refactor)\b",
            "monitoring": r"\b(monitor|alert|metric|log|trace|observe)\b",
            "security": r"\b(security|scan|vulnerability|audit|compliance|pentest)\b",
            "testing": r"\b(test|qa|verify|validate|check|e2e)\b",
            "data": r"\b(data|etl|pipeline|analytics|warehouse)\b",
            "documentation": r"\b(document|doc|readme|wiki|guide)\b",
            "troubleshooting": r"\b(debug|fix|troubleshoot|diagnose|issue)\b",
        }

        task_category = "general"
        for category, pattern in category_patterns.items():
            if re.search(pattern, prompt_lower):
                task_category = category
                break

        # Estimate complexity
        complexity_indicators = len(re.findall(
            r"\b(complex|multiple|all|entire|comprehensive|distributed|microservice)\b",
            prompt_lower
        ))

        if len(prompt) > 200 or complexity_indicators > 2:
            complexity = "high"
        elif len(prompt) > 100 or complexity_indicators > 0:
            complexity = "medium"
        else:
            complexity = "low"

        # Extract action keywords
        action_keywords = re.findall(
            r"\b(deploy|create|build|setup|configure|monitor|test|scan|analyze|fix|update|migrate|debug)\b",
            prompt_lower
        )

        return {
            "task_category": task_category,
            "complexity": complexity,
            "keywords": list(set(action_keywords)),
            "prompt_length": len(prompt),
        }

    def _score_entities(
        self,
        task_analysis: Dict[str, Any],
        entities: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Score and rank entities based on task requirements.

        Scoring:
        - Name/description keyword match: +40 points
        - System prompt relevance: +30 points
        - Tag match: +20 points
        - Complexity match: +10 points

        Returns:
            Sorted list of scored entities (highest first)
        """
        scores = []
        keywords = task_analysis.get("keywords", [])
        task_category = task_analysis.get("task_category", "")

        for entity in entities:
            score = 0
            reasons = []

            # 1. Name/description keyword match (40 points max)
            text_to_search = f"{entity['name']} {entity['description']}".lower()
            keyword_matches = [kw for kw in keywords if kw in text_to_search]
            if keyword_matches:
                keyword_score = min(40, len(keyword_matches) * 13)
                score += keyword_score
                reasons.append(f"Name/desc matches: {', '.join(keyword_matches)}")

            # Also check task category match
            if task_category in text_to_search:
                score += 15
                reasons.append(f"Category match: {task_category}")

            # 2. System prompt relevance (30 points max)
            system_prompt = entity.get("system_prompt", "").lower()
            if system_prompt:
                prompt_matches = [kw for kw in keywords if kw in system_prompt]
                if prompt_matches:
                    prompt_score = min(30, len(prompt_matches) * 10)
                    score += prompt_score
                    reasons.append(f"System prompt matches: {', '.join(prompt_matches)}")

            # 3. Tag match (20 points max)
            entity_tags = set(str(tag).lower() for tag in entity.get("tags", []))
            keyword_set = set(keywords)
            tag_overlap = entity_tags & keyword_set
            if tag_overlap:
                tag_score = min(20, len(tag_overlap) * 10)
                score += tag_score
                reasons.append(f"Tag matches: {', '.join(tag_overlap)}")

            # 4. Complexity match (10 points)
            # High complexity tasks prefer teams over single agents
            if task_analysis.get("complexity") == "high":
                if entity.get("subagent_count", 0) > 0:
                    score += 10
                    reasons.append("Team for complex task")

            # Normalize score to 0-100
            normalized_score = min(100, score)

            # Determine confidence
            if normalized_score >= 70:
                confidence = "high"
            elif normalized_score >= 40:
                confidence = "medium"
            else:
                confidence = "low"

            scores.append({
                "entity_id": entity["id"],
                "entity_name": entity["name"],
                "score": normalized_score,
                "confidence": confidence,
                "reasoning": "; ".join(reasons) if reasons else "Default selection",
            })

        # Sort by score descending
        scores.sort(key=lambda x: x["score"], reverse=True)

        return scores

    async def _llm_recommend_entity_with_retry(
        self,
        prompt: str,
        task_analysis: Dict[str, Any],
        top_entities: List[Dict[str, Any]],
        entity_type: str,
        max_retries: int = 3,
        initial_timeout: float = 10.0
    ) -> Dict[str, Any]:
        """
        Use LLM to make final recommendation with exponential backoff retry.

        Args:
            prompt: Original task prompt
            task_analysis: Pattern-based analysis
            top_entities: Top 3 scored entities
            entity_type: "agent" or "team"
            max_retries: Maximum retry attempts
            initial_timeout: Initial timeout in seconds

        Returns:
            {
                "entity_id": str,
                "reasoning": str
            }
        """
        for attempt in range(max_retries):
            timeout = initial_timeout * (2 ** attempt)

            try:
                logger.debug(
                    "LLM recommendation attempt",
                    attempt=attempt + 1,
                    max_retries=max_retries,
                    timeout=timeout
                )

                result = await asyncio.wait_for(
                    self._llm_recommend_entity(prompt, task_analysis, top_entities, entity_type),
                    timeout=timeout
                )

                return result

            except asyncio.TimeoutError:
                logger.warning(
                    "LLM call timeout",
                    attempt=attempt + 1,
                    timeout=timeout
                )

                if attempt < max_retries - 1:
                    await asyncio.sleep(0.5 * (2 ** attempt))
                else:
                    self._metrics["llm_failures"] += 1
                    raise

            except Exception as e:
                logger.warning(
                    "LLM call failed",
                    attempt=attempt + 1,
                    error=str(e)
                )

                if attempt < max_retries - 1:
                    await asyncio.sleep(0.5 * (2 ** attempt))
                else:
                    self._metrics["llm_failures"] += 1
                    raise

        raise RuntimeError("LLM recommendation exhausted all retries")

    async def _llm_recommend_entity(
        self,
        prompt: str,
        task_analysis: Dict[str, Any],
        top_entities: List[Dict[str, Any]],
        entity_type: str
    ) -> Dict[str, Any]:
        """
        Use LLM to make final recommendation when scores are close.

        Args:
            prompt: Original task prompt
            task_analysis: Pattern-based analysis
            top_entities: Top 3 scored entities
            entity_type: "agent" or "team"

        Returns:
            {
                "entity_id": str,
                "reasoning": str
            }
        """
        self._metrics["llm_calls"] += 1

        # Build context for LLM
        entity_descriptions = []
        for idx, entity in enumerate(top_entities, 1):
            entity_descriptions.append(
                f"{idx}. {entity['entity_name']} (ID: {entity['entity_id']})\n"
                f"   - Score: {entity['score']}\n"
                f"   - Pattern reasoning: {entity['reasoning']}"
            )

        system_prompt = f"""You are an expert at selecting optimal {entity_type}s for tasks.
Analyze the task and recommend the best {entity_type} from the candidates provided.

Consider:
- Task requirements and complexity
- {entity_type.capitalize()} capabilities and descriptions
- Pattern-based analysis results

Respond with ONLY a JSON object (no markdown, no explanation):
{{
  "entity_id": "uuid-of-selected-{entity_type}",
  "reasoning": "Brief explanation (1-2 sentences) of why this {entity_type} is best"
}}"""

        user_prompt = f"""Task prompt: "{prompt}"

Task analysis:
- Category: {task_analysis.get('task_category')}
- Complexity: {task_analysis.get('complexity')}
- Keywords: {', '.join(task_analysis.get('keywords', []))}

Top {entity_type} candidates:
{chr(10).join(entity_descriptions)}

Which {entity_type} should execute this task?"""

        try:
            # Use agno_service for LLM call
            result = await agno_service.execute_agent_async(
                prompt=user_prompt,
                system_prompt=system_prompt,
                model="kubiya/claude-sonnet-4",
                temperature=0.3,
                max_tokens=300,
                stream=False,
            )

            if result.get("success"):
                response_text = result.get("response", "")

                # Parse JSON response
                try:
                    json_match = re.search(r'\{[^}]+\}', response_text, re.DOTALL)
                    if json_match:
                        recommendation = json.loads(json_match.group())

                        if "entity_id" not in recommendation:
                            raise ValueError("LLM response missing entity_id field")

                        logger.info(
                            "LLM entity recommendation success",
                            entity_id=recommendation.get("entity_id"),
                            reasoning=recommendation.get("reasoning", "")[:100]
                        )

                        return recommendation

                    else:
                        logger.warning("No JSON found in LLM response", response=response_text[:200])
                        raise ValueError("Invalid LLM response format")

                except json.JSONDecodeError as e:
                    logger.warning("Failed to parse LLM JSON response", response=response_text[:200], error=str(e))
                    raise
            else:
                error_msg = result.get("error", "Unknown error")
                logger.error("LLM execution failed", error=error_msg)
                raise RuntimeError(f"LLM execution failed: {error_msg}")

        except Exception as e:
            logger.error("LLM recommendation error", error=str(e))
            raise

        # Fallback: return first entity (highest scored)
        logger.info("Using fallback (highest scored entity)")
        return {
            "entity_id": top_entities[0]["entity_id"],
            "reasoning": "Pattern-based selection (LLM unavailable)"
        }

    def _cleanup_cache(self):
        """Remove expired cache entries"""
        current_time = time.time()
        expired_keys = [
            key for key, (_, cached_time) in self._entity_cache.items()
            if current_time - cached_time > self._entity_cache_ttl * 2
        ]

        for key in expired_keys:
            del self._entity_cache[key]

        if expired_keys:
            logger.debug("Cleaned entity cache entries", count=len(expired_keys))

    def get_metrics(self) -> Dict[str, Any]:
        """Get performance metrics for monitoring"""
        return {
            **self._metrics,
            "cache_size": len(self._entity_cache),
        }


# Singleton instance
entity_resolution_service = EntityResolutionService()
