"""
bank profile utilities for personality and background management.
"""

import json
import logging
import re
from typing import Dict, Optional, TypedDict
from pydantic import BaseModel, Field
from ..db_utils import acquire_with_retry
from ..response_models import PersonalityTraits

logger = logging.getLogger(__name__)

DEFAULT_PERSONALITY = {
    "openness": 0.5,
    "conscientiousness": 0.5,
    "extraversion": 0.5,
    "agreeableness": 0.5,
    "neuroticism": 0.5,
    "bias_strength": 0.5,
}


class BankProfile(TypedDict):
    """Type for bank profile data."""
    name: str
    personality: PersonalityTraits
    background: str


class BackgroundMergeResponse(BaseModel):
    """LLM response for background merge with personality inference."""
    background: str = Field(description="Merged background in first person perspective")
    personality: PersonalityTraits = Field(description="Inferred Big Five personality traits")


async def get_bank_profile(pool, bank_id: str) -> BankProfile:
    """
    Get bank profile (name, personality + background).
    Auto-creates bank with default values if not exists.

    Args:
        pool: Database connection pool
        bank_id: bank IDentifier

    Returns:
        BankProfile with name, typed PersonalityTraits, and background
    """
    async with acquire_with_retry(pool) as conn:
        # Try to get existing bank
        row = await conn.fetchrow(
            """
            SELECT name, personality, background
            FROM banks WHERE bank_id = $1
            """,
            bank_id
        )

        if row:
            # asyncpg returns JSONB as a string, so parse it
            personality_data = row["personality"]
            if isinstance(personality_data, str):
                personality_data = json.loads(personality_data)

            return BankProfile(
                name=row["name"],
                personality=PersonalityTraits(**personality_data),
                background=row["background"]
            )

        # Bank doesn't exist, create with defaults
        await conn.execute(
            """
            INSERT INTO banks (bank_id, name, personality, background)
            VALUES ($1, $2, $3::jsonb, $4)
            ON CONFLICT (bank_id) DO NOTHING
            """,
            bank_id,
            bank_id,  # Default name is the bank_id
            json.dumps(DEFAULT_PERSONALITY),
            ""
        )

        return BankProfile(
            name=bank_id,
            personality=PersonalityTraits(**DEFAULT_PERSONALITY),
            background=""
        )


async def update_bank_personality(
    pool,
    bank_id: str,
    personality: Dict[str, float]
) -> None:
    """
    Update bank personality traits.

    Args:
        pool: Database connection pool
        bank_id: bank IDentifier
        personality: Dict with Big Five traits + bias_strength (all 0-1)
    """
    # Ensure bank exists first
    await get_bank_profile(pool, bank_id)

    async with acquire_with_retry(pool) as conn:
        await conn.execute(
            """
            UPDATE banks
            SET personality = $2::jsonb,
                updated_at = NOW()
            WHERE bank_id = $1
            """,
            bank_id,
            json.dumps(personality)
        )


async def merge_bank_background(
    pool,
    llm_config,
    bank_id: str,
    new_info: str,
    update_personality: bool = True
) -> dict:
    """
    Merge new background information with existing background using LLM.
    Normalizes to first person ("I") and resolves conflicts.
    Optionally infers personality traits from the merged background.

    Args:
        pool: Database connection pool
        llm_config: LLM configuration for background merging
        bank_id: bank IDentifier
        new_info: New background information to add/merge
        update_personality: If True, infer Big Five traits from background (default: True)

    Returns:
        Dict with 'background' (str) and optionally 'personality' (dict) keys
    """
    # Get current profile
    profile = await get_bank_profile(pool, bank_id)
    current_background = profile["background"]

    # Use LLM to merge backgrounds and optionally infer personality
    result = await _llm_merge_background(
        llm_config,
        current_background,
        new_info,
        infer_personality=update_personality
    )

    merged_background = result["background"]
    inferred_personality = result.get("personality")

    # Update in database
    async with acquire_with_retry(pool) as conn:
        if inferred_personality:
            # Update both background and personality
            await conn.execute(
                """
                UPDATE banks
                SET background = $2,
                    personality = $3::jsonb,
                    updated_at = NOW()
                WHERE bank_id = $1
                """,
                bank_id,
                merged_background,
                json.dumps(inferred_personality)
            )
        else:
            # Update only background
            await conn.execute(
                """
                UPDATE banks
                SET background = $2,
                    updated_at = NOW()
                WHERE bank_id = $1
                """,
                bank_id,
                merged_background
            )

    response = {"background": merged_background}
    if inferred_personality:
        response["personality"] = inferred_personality

    return response


async def _llm_merge_background(
    llm_config,
    current: str,
    new_info: str,
    infer_personality: bool = False
) -> dict:
    """
    Use LLM to intelligently merge background information.
    Optionally infer Big Five personality traits from the merged background.

    Args:
        llm_config: LLM configuration to use
        current: Current background text
        new_info: New information to merge
        infer_personality: If True, also infer personality traits

    Returns:
        Dict with 'background' (str) and optionally 'personality' (dict) keys
    """
    if infer_personality:
        prompt = f"""You are helping maintain a memory bank's background/profile and infer their personality. You MUST respond with ONLY valid JSON.

Current background: {current if current else "(empty)"}

New information to add: {new_info}

Instructions:
1. Merge the new information with the current background
2. If there are conflicts (e.g., different birthplaces), the NEW information overwrites the old
3. Keep additions that don't conflict
4. Output in FIRST PERSON ("I") perspective
5. Be concise - keep merged background under 500 characters
6. Infer Big Five personality traits from the merged background:
   - Openness: 0.0-1.0 (creativity, curiosity, openness to new ideas)
   - Conscientiousness: 0.0-1.0 (organization, discipline, goal-directed)
   - Extraversion: 0.0-1.0 (sociability, assertiveness, energy from others)
   - Agreeableness: 0.0-1.0 (cooperation, empathy, consideration)
   - Neuroticism: 0.0-1.0 (emotional sensitivity, anxiety, stress response)
   - Bias Strength: 0.0-1.0 (how much personality influences opinions)

CRITICAL: You MUST respond with ONLY a valid JSON object. No markdown, no code blocks, no explanations. Just the JSON.

Format:
{{
  "background": "the merged background text in first person",
  "personality": {{
    "openness": 0.7,
    "conscientiousness": 0.6,
    "extraversion": 0.5,
    "agreeableness": 0.8,
    "neuroticism": 0.4,
    "bias_strength": 0.6
  }}
}}

Trait inference examples:
- "creative artist" → openness: 0.8+, bias_strength: 0.6
- "organized engineer" → conscientiousness: 0.8+, openness: 0.5-0.6
- "startup founder" → openness: 0.8+, extraversion: 0.7+, neuroticism: 0.3-0.4
- "risk-averse analyst" → openness: 0.3-0.4, conscientiousness: 0.8+, neuroticism: 0.6+
- "rational and diligent" → conscientiousness: 0.7+, openness: 0.6+
- "passionate and dramatic" → extraversion: 0.7+, neuroticism: 0.6+, openness: 0.7+"""
    else:
        prompt = f"""You are helping maintain a memory bank's background/profile.

Current background: {current if current else "(empty)"}

New information to add: {new_info}

Instructions:
1. Merge the new information with the current background
2. If there are conflicts (e.g., different birthplaces), the NEW information overwrites the old
3. Keep additions that don't conflict
4. Output in FIRST PERSON ("I") perspective
5. Be concise - keep it under 500 characters
6. Return ONLY the merged background text, no explanations

Merged background:"""

    try:
        # Prepare messages
        messages = [{"role": "user", "content": prompt}]

        if infer_personality:
            # Use structured output with Pydantic model for personality inference
            try:
                parsed = await llm_config.call(
                    messages=messages,
                    response_format=BackgroundMergeResponse,
                    scope="bank_background",
                    temperature=0.3,
                    max_tokens=8192
                )
                logger.info(f"Successfully got structured response: background={parsed.background[:100]}")

                # Convert Pydantic model to dict format
                return {
                    "background": parsed.background,
                    "personality": parsed.personality.model_dump()
                }
            except Exception as e:
                logger.warning(f"Structured output failed, falling back to manual parsing: {e}")
                # Fall through to manual parsing below

        # Manual parsing fallback or non-personality merge
        content = await llm_config.call(
            messages=messages,
            scope="bank_background",
            temperature=0.3,
            max_tokens=8192
        )

        logger.info(f"LLM response for background merge (first 500 chars): {content[:500]}")

        if infer_personality:
            # Parse JSON response - try multiple extraction methods
            result = None

            # Method 1: Direct parse
            try:
                result = json.loads(content)
                logger.info("Successfully parsed JSON directly")
            except json.JSONDecodeError:
                pass

            # Method 2: Extract from markdown code blocks
            if result is None:
                # Remove markdown code blocks
                code_block_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
                if code_block_match:
                    try:
                        result = json.loads(code_block_match.group(1))
                        logger.info("Successfully extracted JSON from markdown code block")
                    except json.JSONDecodeError:
                        pass

            # Method 3: Find nested JSON structure
            if result is None:
                # Look for JSON object with nested structure
                json_match = re.search(r'\{[^{}]*"background"[^{}]*"personality"[^{}]*\{[^{}]*\}[^{}]*\}', content, re.DOTALL)
                if json_match:
                    try:
                        result = json.loads(json_match.group())
                        logger.info("Successfully extracted JSON using nested pattern")
                    except json.JSONDecodeError:
                        pass

            # All parsing methods failed - use fallback
            if result is None:
                logger.warning(f"Failed to extract JSON from LLM response. Raw content: {content[:200]}")
                # Fallback: use new_info as background with default personality
                return {
                    "background": new_info if new_info else current if current else "",
                    "personality": DEFAULT_PERSONALITY.copy()
                }

            # Validate personality values
            personality = result.get("personality", {})
            for key in ["openness", "conscientiousness", "extraversion",
                       "agreeableness", "neuroticism", "bias_strength"]:
                if key not in personality:
                    personality[key] = 0.5  # Default to neutral
                else:
                    # Clamp to [0, 1]
                    personality[key] = max(0.0, min(1.0, float(personality[key])))

            result["personality"] = personality

            # Ensure background exists
            if "background" not in result or not result["background"]:
                result["background"] = new_info if new_info else ""

            return result
        else:
            # Just background merge
            merged = content
            if not merged or merged.lower() in ["(empty)", "none", "n/a"]:
                merged = new_info if new_info else ""
            return {"background": merged}

    except Exception as e:
        logger.error(f"Error merging background with LLM: {e}")
        # Fallback: just append new info
        if current:
            merged = f"{current} {new_info}".strip()
        else:
            merged = new_info

        result = {"background": merged}
        if infer_personality:
            result["personality"] = DEFAULT_PERSONALITY.copy()
        return result


async def list_banks(pool) -> list:
    """
    List all banks in the system.

    Args:
        pool: Database connection pool

    Returns:
        List of dicts with bank_id, name, personality, background, created_at, updated_at
    """
    async with acquire_with_retry(pool) as conn:
        rows = await conn.fetch(
            """
            SELECT bank_id, name, personality, background, created_at, updated_at
            FROM banks
            ORDER BY updated_at DESC
            """
        )

        result = []
        for row in rows:
            # asyncpg returns JSONB as a string, so parse it
            personality_data = row["personality"]
            if isinstance(personality_data, str):
                personality_data = json.loads(personality_data)

            result.append({
                "bank_id": row["bank_id"],
                "name": row["name"],
                "personality": personality_data,
                "background": row["background"],
                "created_at": row["created_at"].isoformat() if row["created_at"] else None,
                "updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
            })

        return result