src/pinecone_extractor.py

import os
import logging
import re
from typing import Dict, Optional

from pinecone import Pinecone
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)


class PineconeUserContextExtractor:
    """
    Extracts structured user context from Pinecone.

    Design goals:
    - Stable contract for SCN
    - Extensible for future user attributes
    - Safe against missing metadata
    - No hardcoded vector dimension
    """

    def __init__(self, index_name: Optional[str] = None):
        self.index_name = index_name or os.getenv("PINECONE_INDEX_NAME")
        self.api_key = os.getenv("PINECONE_API_KEY")

        if not self.api_key:
            raise ValueError("PINECONE_API_KEY not set")

        if not self.index_name:
            raise ValueError("PINECONE_INDEX_NAME not set")

        self.pc = Pinecone(api_key=self.api_key)
        self.index = self.pc.Index(self.index_name)

        # 🔥 Read actual index dimension dynamically
        stats = self.index.describe_index_stats()
        self.dimension = stats.get("dimension")

        if not self.dimension:
            raise RuntimeError("Could not determine Pinecone index dimension")

        logger.info(
            f"Pinecone extractor initialized | index={self.index_name}, dimension={self.dimension}"
        )

    # -------------------------------------------------------
    # Public API
    # -------------------------------------------------------

    def get_user_context(self, namespace: str) -> Dict:
        """
        Returns structured user context.

        Output contract (stable):
        {
            user_name: str | None,
            role: str | None,
            age: int | None,
            summary: str | None,
            resume_summary: str | None,
            skills: list[str],
            interests: list[str],
            traits: list[str],
        }
        """

        context = {
            "user_name": None,
            "role": None,
            "age": None,
            "summary": None,
            "resume_summary": None,
            "skills": [],
            "interests": [],
            "traits": [],
        }

        try:
            dummy_vector = [0.0] * self.dimension

            results = self.index.query(
                vector=dummy_vector,
                top_k=20,
                include_metadata=True,
                namespace=str(namespace),
            )

            matches = results.get("matches", []) or []

            for match in matches:
                meta = match.get("metadata") or {}
                doc_type = meta.get("doc_type")

                # -------------------------------------------
                # Conversation Summary (Highest Signal)
                # -------------------------------------------
                if doc_type == "conversation_summary":
                    if not context["summary"]:
                        context["summary"] = meta.get("text")

                # -------------------------------------------
                # Resume Summary
                # -------------------------------------------
                elif doc_type == "resume_summary":
                    text = meta.get("text", "")

                    if not context["resume_summary"]:
                        context["resume_summary"] = text

                    # Extract structured fields from text if needed
                    self._extract_from_resume_text(text, context)

                # -------------------------------------------
                # Optional structured metadata types
                # -------------------------------------------
                elif doc_type == "skills":
                    context["skills"].extend(meta.get("items", []))

                elif doc_type == "interests":
                    context["interests"].extend(meta.get("items", []))

                elif doc_type == "traits":
                    context["traits"].extend(meta.get("items", []))

            # Deduplicate lists
            for key in ["skills", "interests", "traits"]:
                context[key] = list(set(context[key]))

            logger.info(
                f"Pinecone context extracted | namespace={namespace} | "
                f"summary={'YES' if context['summary'] else 'NO'} | "
                f"resume={'YES' if context['resume_summary'] else 'NO'}"
            )

        except Exception as e:
            logger.error(
                f"Failed to extract Pinecone context for namespace={namespace}: {e}",
                exc_info=True
            )

        return context

    # -------------------------------------------------------
    # Internal Helpers
    # -------------------------------------------------------

    def _extract_from_resume_text(self, text: str, context: Dict) -> None:
        if not text:
            return

        # ---------------------------
        # Extract Name
        # ---------------------------
        if not context["user_name"]:
            name_match = re.search(r"Name:\s*(.+)", text)
            if name_match:
                context["user_name"] = name_match.group(1).strip()

        # ---------------------------
        # Extract Role
        # ---------------------------
        if not context["role"]:
            role_match = re.search(r"Role:\s*(.+)", text)
            if role_match:
                context["role"] = role_match.group(1).strip()

        # ---------------------------
        # Extract Age
        # ---------------------------
        if not context["age"]:
            age_match = re.search(r"Age:\s*(\d+)", text)
            if age_match:
                context["age"] = int(age_match.group(1))

        # ---------------------------
        # Extract Skills Block
        # ---------------------------
        skills_match = re.search(r"Skills:\s*(.*?)\n\n", text, re.DOTALL)
        if skills_match:
            skills_block = skills_match.group(1)

            skill_lines = re.findall(r"-\s*(.+)", skills_block)

            parsed_skills = []
            for line in skill_lines:
                parts = [s.strip() for s in line.split(",")]
                parsed_skills.extend(parts)

            context["skills"].extend(parsed_skills)

        # ---------------------------
        # Extract Learning Goals
        # ---------------------------
        goals_match = re.search(r"Learning Goals:\s*(.*?)\n\n", text, re.DOTALL)
        if goals_match:
            goals_block = goals_match.group(1)
            goals = re.findall(r"-\s*(.+)", goals_block)
            context["interests"].extend(goals)

        # ---------------------------
        # Extract Hobbies
        # ---------------------------
        hobby_match = re.search(r"Hobbies:\s*(.*?)$", text, re.DOTALL)
        if hobby_match:
            hobby_block = hobby_match.group(1)
            hobbies = re.findall(r"-\s*(.+)", hobby_block)
            context["interests"].extend(hobbies)

    def get_latest_roadmap(self, user_id: str) -> Optional[dict]:
        import json

        try:
            dummy_vector = [0.0] * self.dimension

            response = self.index.query(
                vector=dummy_vector,
                top_k=10,
                namespace=str(user_id),
                include_metadata=True,
                filter={"doc_type": {"$eq": "roadmap"}}
            )

            matches = response.get("matches", []) or []

            if not matches:
                logger.warning(f"[ROADMAP FETCH] No roadmap found for user: {user_id}")
                return None

            matches.sort(
                key=lambda x: (x.get("metadata") or {}).get("generated_at", ""),
                reverse=True
            )

            best_metadata = matches[0].get("metadata") or {}

            if not best_metadata.get("full_roadmap_stored"):
                logger.warning(f"[ROADMAP FETCH] full_roadmap_stored=False for user: {user_id}")
                return None

            raw_json = best_metadata.get("full_roadmap_json", "")
            if not raw_json:
                return None

            roadmap = json.loads(raw_json)
            roadmap = self._normalize_roadmap(roadmap)

            logger.info(
                f"[ROADMAP FETCH] ✓ Roadmap retrieved for user {user_id} | "
                f"milestones={len(roadmap.get('milestones', []))} | "
                f"target_role={roadmap.get('target_role', 'unknown')}"
            )
            return roadmap

        except Exception as e:
            logger.error(f"[ROADMAP FETCH] Error: {e}", exc_info=True)
            return None


    def _normalize_roadmap(self, roadmap: dict) -> dict:
        """
        Normalizes abbreviated milestone keys to standard names.
        Handles both old schema (abbreviated) and new schema (full names).
        """
        for milestone in roadmap.get("milestones", []):

            # Normalize milestone-level fields
            if not milestone.get("identity_label"):
                milestone["identity_label"] = (
                    milestone.get("t") or
                    milestone.get("label") or
                    milestone.get("milestone_id", "")
                )

            if not milestone.get("market_value_display"):
                milestone["market_value_display"] = (
                    milestone.get("sal") or
                    milestone.get("salary") or
                    ""
                )

            if not milestone.get("identity_statement"):
                milestone["identity_statement"] = (
                    milestone.get("o") or
                    milestone.get("outcome") or
                    ""
                )

            # Normalize module-level fields
            for module in milestone.get("modules", []):
                if not module.get("module_id"):
                    module["module_id"] = module.get("id", "")

                # Normalize skill-level fields
                for skill in module.get("skills", []):
                    if not skill.get("description"):
                        skill["description"] = skill.get("n", "")

                    # Ensure why_this_skill exists (needed for Beat 5)
                    if not skill.get("why_this_skill"):
                        scenario = skill.get("content_flow", {}).get("scenario", {})
                        skill["why_this_skill"] = (
                            f"This skill is tested in: {scenario.get('title', '')}"
                            if scenario.get("title") else ""
                        )

        return roadmap

    def get_onboarding_summary(self, user_id: str) -> Optional[dict]:
        """
        Fetch structured onboarding summary stored by seed script or memory_service.
        Returns the parsed summary dict or None.
        """
        import json
        try:
            vector_id = f"{user_id}_onboarding_summary"
            fetch_result = self.index.fetch(
                ids=[vector_id],
                namespace=str(user_id)
            )
            if not (fetch_result and fetch_result.vectors and vector_id in fetch_result.vectors):
                logger.warning(f"[SUMMARY] No onboarding_summary vector for: {user_id}")
                return None

            metadata = fetch_result.vectors[vector_id].metadata or {}
            text = metadata.get("text", "")
            if not text:
                return None

            data = json.loads(text)
            logger.info(f"[SUMMARY] Found onboarding_summary for: {user_id}")
            return data

        except Exception as e:
            logger.error(f"[SUMMARY] Error fetching onboarding_summary for {user_id}: {e}")
            return None