"""
FoundryGraph AI Screening Service.

Uses Vertex AI (Gemini) to:
1. Score data quality (0-100)
2. Identify issues and conflicts
3. Resolve multi-source conflicts
4. Infer parent/subsidiary relationships
5. Mine aliases and name variants
"""

import json
import logging
import os
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

import vertexai
from vertexai.generative_models import GenerationConfig, GenerativeModel

logger = logging.getLogger(__name__)

MODEL_ID = os.getenv("FM_FOUNDRYGRAPH_MODEL_ID", "gemini-2.0-flash-lite-001")
# Keep a simple version string for tracking
MODEL_VERSION = "v1.0"

def _vertex_location_for_model(model: Optional[str], location: str) -> str:
    if (model or "").strip().lower().startswith("gemini-3-"):
        return "global"
    return location


@dataclass
class QualityAssessment:
    """Result of AI quality screening."""

    score: int
    issues: List[str]
    suggestions: List[str]
    conflicts: List[Dict[str, Any]]
    promotion_recommendation: str  # promote | review | reject
    confidence: float
    model_version: str
    screened_at: str


@dataclass
class AliasResult:
    """Mined alias/name variant."""

    alias: str
    alias_type: str  # abbreviation, former_name, dba, translation
    confidence: float
    evidence: str


@dataclass
class ParentInference:
    """Inferred parent-child relationship."""

    parent_wikidata_id: Optional[str]
    parent_name: str
    relationship_type: str  # parent, subsidiary, division, brand
    confidence: float
    evidence: List[str]


class FoundryGraphAIScreener:
    """
    AI screening service for FoundryGraph data quality.

    Philosophy:
    - AI creates durable assets (quality scores, aliases, relationships)
    - NOT one-off answers
    """

    def __init__(self, project_id: str, location: Optional[str] = None):
        base_location = location or os.getenv("FM_VERTEX_AI_LOCATION", "us-central1")
        resolved_location = _vertex_location_for_model(MODEL_ID, base_location)
        if resolved_location != base_location:
            logger.info("Vertex AI: overriding location to global for model %s", MODEL_ID)
        vertexai.init(project=project_id, location=resolved_location)
        self.model = GenerativeModel(MODEL_ID)
        self.generation_config = GenerationConfig(
            temperature=0.1,
            max_output_tokens=2048,
            response_mime_type="application/json",
        )

    async def screen_company(self, company: Dict[str, Any]) -> QualityAssessment:
        """Screen a single company record for quality."""
        prompt = f"""You are a data quality analyst for a B2B company database.

Analyze this company record and provide a quality assessment.

COMPANY RECORD:
{json.dumps(company, indent=2, default=str)}

EVALUATION CRITERIA:
1. COMPLETENESS (0-25 points):
   - Has domain/website: +10
   - Has industry: +5
   - Has headquarters: +5
   - Has external ID (ticker, LEI): +5

2. CONSISTENCY (0-25 points):
   - Domain matches company name pattern: +10
   - Website URL is valid format: +5
   - Name doesn't look like spam/placeholder: +10

3. PLAUSIBILITY (0-25 points):
   - Revenue reasonable for company type: +10
   - Employee count reasonable: +10
   - Inception date plausible: +5

4. UNIQUENESS (0-25 points):
   - Not obviously a duplicate of another entity: +15
   - Name is distinctive enough for matching: +10

Respond with JSON:
{{
  "score": <0-100>,
  "issues": ["list of specific issues found"],
  "suggestions": ["actionable fixes"],
  "conflicts": [{{"field": "...", "values": [...], "resolution": "..."}}],
  "promotion_recommendation": "promote|review|reject",
  "confidence": <0.0-1.0>,
  "reasoning": "brief explanation"
}}
"""
        try:
            response = await self.model.generate_content_async(prompt, generation_config=self.generation_config)
            result = json.loads(response.text)
            return QualityAssessment(
                score=int(result.get("score", 0)),
                issues=result.get("issues", []),
                suggestions=result.get("suggestions", []),
                conflicts=result.get("conflicts", []),
                promotion_recommendation=result.get("promotion_recommendation", "review"),
                confidence=float(result.get("confidence", 0.5)),
                model_version=MODEL_VERSION,
                screened_at=datetime.now(timezone.utc).isoformat(),
            )
        except Exception as exc:  # noqa: BLE001
            logger.error("AI screening failed for %s: %s", company.get("wikidata_id"), exc)
            return QualityAssessment(
                score=50,
                issues=["AI screening failed - manual review needed"],
                suggestions=[],
                conflicts=[],
                promotion_recommendation="review",
                confidence=0.0,
                model_version=MODEL_VERSION,
                screened_at=datetime.now(timezone.utc).isoformat(),
            )

    async def mine_aliases(self, company: Dict[str, Any]) -> List[AliasResult]:
        """
        Mine aliases and name variants from company data.

        Per Non-Negotiable #3: ONLY allow aliases derivable from the record.
        """
        prompt = f"""You are a B2B data expert. Analyze this company and identify name variants.

COMPANY:
Name: {company.get('label')}
Domain: {company.get('domain')}
Industry: {company.get('industries')}

STRICT RULES (Non-Negotiable #3 - Hallucination Prevention):
1. ONLY generate variants DERIVABLE from the provided data
2. ALLOWED types:
   - suffix_variant: "Inc" ↔ "Inc." ↔ "Incorporated" ↔ "" (removing suffix)
   - punctuation_variant: "AT&T" ↔ "AT and T" ↔ "ATT"
   - whitespace_variant: "Wal-Mart" ↔ "Walmart" ↔ "Wal Mart"
   - abbreviation: ONLY if derivable from initials (e.g., "International Business Machines" → "IBM")
3. NOT ALLOWED:
   - Former names you "know" from general knowledge
   - Acquisition names (don't add former parents unless in the record)
   - Brand names not in the record
   - Guesses or speculation

For each alias, provide evidence showing how it derives from the input.

Respond with JSON:
{{
  "aliases": [
    {{
      "alias": "the variant",
      "type": "suffix_variant|punctuation_variant|whitespace_variant|abbreviation",
      "confidence": 0.0-1.0,
      "evidence": "EXACTLY how this derives from the input name",
      "derivation_rule": "specific transformation applied"
    }}
  ]
}}

If no valid aliases can be derived, return {{"aliases": []}}
"""
        try:
            response = await self.model.generate_content_async(prompt, generation_config=self.generation_config)
            result = json.loads(response.text)
            valid_types = {"suffix_variant", "punctuation_variant", "whitespace_variant", "abbreviation"}

            aliases: List[AliasResult] = []
            for alias in result.get("aliases", []):
                if alias.get("type") not in valid_types:
                    continue
                if not alias.get("evidence") or len(alias.get("evidence", "")) < 10:
                    continue
                if alias.get("confidence", 0) < 0.8:
                    continue
                aliases.append(
                    AliasResult(
                        alias=alias["alias"],
                        alias_type=alias.get("type", "variation"),
                        confidence=alias.get("confidence", 0.8),
                        evidence=alias.get("evidence", ""),
                    )
                )
            return aliases
        except Exception as exc:  # noqa: BLE001
            logger.error("Alias mining failed for %s: %s", company.get("wikidata_id"), exc)
            return []

    async def resolve_conflict(self, field: str, sources: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Resolve conflict between multiple sources for a field."""
        prompt = f"""You are resolving a data conflict for a company database.

FIELD: {field}

CONFLICTING SOURCES:
{json.dumps(sources, indent=2, default=str)}

Rules for resolution:
1. SEC filings > Wikidata for US financial data (more authoritative)
2. More recent > older (check dates)
3. Primary source > aggregator
4. If truly ambiguous, flag for manual review

Respond with JSON:
{{
  "resolved_value": "the winning value",
  "source": "which source to trust",
  "confidence": 0.0-1.0,
  "reasoning": "why this source wins",
  "needs_review": true/false
}}
"""
        try:
            response = await self.model.generate_content_async(prompt, generation_config=self.generation_config)
            return json.loads(response.text)
        except Exception as exc:  # noqa: BLE001
            logger.error("Conflict resolution failed for %s: %s", field, exc)
            return {
                "resolved_value": None,
                "source": None,
                "confidence": 0.0,
                "reasoning": "AI resolution failed",
                "needs_review": True,
            }

    async def infer_parent_relationship(
        self, company: Dict[str, Any], candidates: List[Dict[str, Any]]
    ) -> Optional[ParentInference]:
        """Infer parent-subsidiary relationship from company data."""
        candidate_list = []
        for idx, candidate in enumerate(candidates[:10]):
            candidate_list.append(
                {
                    "index": idx,
                    "wikidata_id": candidate.get("wikidata_id"),
                    "name": candidate.get("label") or candidate.get("name"),
                    "domain": candidate.get("domain"),
                }
            )

        prompt = f"""You are analyzing corporate structure for a B2B database.

STRICT RULES (Non-Negotiable #3 - Hallucination Prevention):
1. You MUST choose a parent ONLY from the CANDIDATE LIST below
2. If the true parent is NOT in the candidate list, return has_parent: false
3. NEVER invent parent companies from general knowledge
4. NEVER use acquisition history you "know" but isn't in the data
5. The parent_wikidata_id MUST match exactly one candidate's wikidata_id

COMPANY TO ANALYZE:
{{
  "wikidata_id": "{company.get('wikidata_id')}",
  "name": "{company.get('label') or company.get('name')}",
  "domain": "{company.get('domain')}",
  "industries": {json.dumps(company.get('industries', []))},
  "headquarters": "{company.get('headquarters', {}) if isinstance(company.get('headquarters'), dict) else company.get('headquarters')}"
}}

CANDIDATE PARENT LIST (choose ONLY from these or return null):
{json.dumps(candidate_list, indent=2)}

ALLOWED SIGNALS (from the provided data ONLY):
1. Domain family: company domain is subdomain/variant of candidate domain
2. Name containment: company name contains candidate name exactly
3. Shared industry + headquarters (weak signal, needs other evidence)

NOT ALLOWED:
- Inventing parent companies not in the candidate list
- Guessing based on similar names without exact containment

Respond with JSON:
{{
  "has_parent": true/false,
  "parent_index": 0-9 or null (index from candidate list),
  "parent_wikidata_id": "Q123 from candidate list or null",
  "relationship_type": "subsidiary|division|brand|null",
  "confidence": 0.0-1.0,
  "evidence": ["ONLY cite data from the provided records"]
}}
"""
        try:
            response = await self.model.generate_content_async(prompt, generation_config=self.generation_config)
            result = json.loads(response.text)
            if not result.get("has_parent") or result.get("confidence", 0) < 0.7:
                return None

            claimed_parent_id = result.get("parent_wikidata_id")
            valid_parent_ids = {c.get("wikidata_id") for c in candidates[:10]}
            if claimed_parent_id not in valid_parent_ids:
                logger.warning(
                    "AI hallucinated parent %s for %s - not in candidate list. Rejecting.",
                    claimed_parent_id,
                    company.get("wikidata_id"),
                )
                return None

            parent_name = ""
            for candidate in candidates[:10]:
                if candidate.get("wikidata_id") == claimed_parent_id:
                    parent_name = candidate.get("label") or candidate.get("name") or ""
                    break

            return ParentInference(
                parent_wikidata_id=claimed_parent_id,
                parent_name=parent_name,
                relationship_type=result.get("relationship_type", "subsidiary"),
                confidence=result.get("confidence", 0.7),
                evidence=result.get("evidence", []),
            )
        except Exception as exc:  # noqa: BLE001
            logger.error("Parent inference failed for %s: %s", company.get("wikidata_id"), exc)
            return None


async def screen_companies_batch(
    screener: FoundryGraphAIScreener, companies: List[Dict[str, Any]], batch_size: int = 10
) -> List[QualityAssessment]:
    """Screen multiple companies with batching for efficiency."""
    import asyncio

    results: List[QualityAssessment] = []
    for idx in range(0, len(companies), batch_size):
        batch = companies[idx : idx + batch_size]
        tasks = [screener.screen_company(company) for company in batch]
        batch_results = await asyncio.gather(*tasks, return_exceptions=True)
        for result in batch_results:
            if isinstance(result, Exception):
                logger.error("Batch screening error: %s", result)
                result = QualityAssessment(
                    score=50,
                    issues=["Screening failed"],
                    suggestions=[],
                    conflicts=[],
                    promotion_recommendation="review",
                    confidence=0.0,
                    model_version=MODEL_VERSION,
                    screened_at=datetime.now(timezone.utc).isoformat(),
                )
            results.append(result)
        await asyncio.sleep(0.5)
    return results
