from __future__ import annotations
from datetime import UTC, datetime
from pydantic import model_validator, Field, BaseModel  # type: ignore
from typing import (
    Dict,
    Literal,
    ClassVar,
    cast,
    Any,
    Tuple,
    AsyncGenerator,
    List,
    Optional
#    override
)  # NOTE: You must import Optional for late binding to work
import inspect
import random
import re
import json
import asyncio
import time
import asyncio
import time
import os
import random
from names_dataset import NameDataset, NameWrapper # type: ignore

from .base import Agent, agent_registry, LLMMessage
from models import ApiActivityType, Candidate, ChatMessage, ChatMessageError, ChatMessageMetaData, ApiMessageType, ChatMessageStatus, ChatMessageStreaming, ChatMessageUser, ChatOptions, ChatSenderType, ApiStatusType, Tunables
import model_cast
from logger import logger
import defines
import backstory_traceback as traceback

seed = int(time.time())
random.seed(seed)

emptyUser = {
    "profile_url": "",
    "description": "",
    "rag_content_size": 0,
    "username": "",
    "first_name": "",
    "last_name": "",
    "full_name": "",
    "email": "",
    "phone": "",
    "title": "",
    "contact_info": {},
    "questions": [],
}

generate_persona_system_prompt = """\
You are a casting director for a movie. Your job is to provide information on ficticious personas for use in a screen play.

All response field MUST BE IN ENGLISH, regardless of ethnicity.

You will be provided with defaults to use if not specified by the user:

```json
{
"age": number,
"gender": "male" | "female",
"ethnicity": string,
"full_name": string,
"first_name": string,
"last_name": string,
}
```

Additional information provided in the user message can override those defaults.

You need to randomly assign an English username (can include numbers), a first name, last name, and a two English sentence description of 
that individual's work given the demographics provided.

Your response must be in JSON. 
Provide only the JSON response, and match the field names EXACTLY. 
Provide all information in English ONLY, with no other commentary:

```json
{
"username": string,  # A likely-to-be unique username, no more than 15 characters (can include numbers and letters but no special characters)
"description": string, # One to two sentence description of their job
"location": string, # In the location, provide ALL of: City, State/Region, and Country
"phone": string, # Location appropriate phone number with area code
"email": string, # primary email address
"title": string, # Job title of their current job
}
```

Make sure to provide a username and that the field name for the job description is "description".

DO NOT infer, imply, abbreviate, or state the ethnicity, gender, or age in the username or description. 
You are providing those only for use later by the system when casting individuals for the role.
"""

generate_resume_system_prompt = """
You are a creative writing casting director. As part of the casting, you are building backstories about individuals. The first part
of that is to create an in-depth resume for the person. You will be provided with the following information:

```json
"full_name": string,   # Person full name
"location": string,    # Location of residence
"age": number,         # Age of candidate
"description": string  # A brief description of the person
```

Use that information to invent a full career resume. Include sections such as:

* Contact information
* Job goal
* Top skills
* Detailed work history. If they are under the age of 25, you might include skills, hobbies, or volunteering they may have done while an adolescent
* In the work history, provide company names, years of employment, and their role
* Education

Provide the resume in Markdown format. DO NOT provide any commentary before or after the resume.
"""

# Debug version to identify the issue
import random
import logging

# Set up logging to see what's happening
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class EthnicNameGenerator:
    def __init__(self):
        try:
            from names_dataset import NameDataset # type: ignore
            self.nd = NameDataset()
        except ImportError:
            logger.error("NameDataset not available. Please install: pip install names-dataset")
            self.nd = None
        except Exception as e:
            logger.error(f"Error initializing NameDataset: {e}")
            self.nd = None
        
        # US Census 2020 approximate ethnic distribution
        self.ethnic_weights = {
            'White': 0.576,
            'Hispanic': 0.186,
            'Black': 0.134,
            'Asian': 0.062,
            'Native American': 0.013,
            'Pacific Islander': 0.003,
            'Mixed/Other': 0.026
        }
        
        # Map ethnicities to countries (using alpha-2 codes that NameDataset uses)
        self.ethnic_country_mapping = {
            'White': ['US', 'GB', 'DE', 'IE', 'IT', 'PL', 'FR', 'CA', 'AU'],
            'Hispanic': ['MX', 'ES', 'CO', 'PE', 'AR', 'CU', 'VE', 'CL'],
            'Black': ['US'],  # African American names
            'Asian': ['CN', 'IN', 'PH', 'VN', 'KR', 'JP', 'TH', 'MY'],
            'Native American': ['US'],
            'Pacific Islander': ['US'],
            'Mixed/Other': ['US']
        }
    
    def get_weighted_ethnicity(self) -> str:
        """Select ethnicity based on US demographic weights"""
        ethnicities = list(self.ethnic_weights.keys())
        weights = list(self.ethnic_weights.values())
        return random.choices(ethnicities, weights=weights)[0]
    
    def get_names_by_criteria(self, countries: List[str], gender: Optional[str] = None, 
                             n: int = 50, use_first_names: bool = True) -> List[str]:
        """Get names matching criteria using NameDataset's get_top_names method"""
        if not self.nd:
            return []
        
        all_names = []
        for country_code in countries:
            try:
                # Get top names for this country
                top_names = self.nd.get_top_names(
                    n=n, 
                    use_first_names=use_first_names,
                    country_alpha2=country_code,
                    gender=gender
                )
                
                if country_code in top_names:
                    if use_first_names and gender:
                        # For first names with gender specified
                        gender_key = 'M' if gender.upper() in ['M', 'MALE'] else 'F'
                        if gender_key in top_names[country_code]:
                            all_names.extend(top_names[country_code][gender_key])
                    elif use_first_names:
                        # For first names without gender (get both)
                        for gender_names in top_names[country_code].values():
                            all_names.extend(gender_names)
                    else:
                        # For last names
                        all_names.extend(top_names[country_code])
                        
            except Exception as e:
                logger.debug(f"Error getting names for {country_code}: {e}")
                continue
        
        return list(set(all_names))  # Remove duplicates
    
    def get_name_by_ethnicity(self, ethnicity: str, gender: str = 'random') -> Tuple[str, str, str, str]:
        """Generate a name based on ethnicity using the correct NameDataset API"""
        if gender == 'random':
            gender = random.choice(['Male', 'Female'])
        
        countries = self.ethnic_country_mapping.get(ethnicity, ['US'])
        
        # Get first names
        first_names = self.get_names_by_criteria(
            countries=countries, 
            gender=gender, 
            use_first_names=True
        )
        
        # Get last names
        last_names = self.get_names_by_criteria(
            countries=countries, 
            use_first_names=False
        )
        
        # Select names or use fallbacks
        if first_names:
            first_name = random.choice(first_names)
        else:
            first_name = self._get_fallback_first_name(gender, ethnicity)
            logger.info(f"Using fallback first name for {ethnicity} {gender}")
        
        if last_names:
            last_name = random.choice(last_names)
        else:
            last_name = self._get_fallback_last_name(ethnicity)
            logger.info(f"Using fallback last name for {ethnicity}")
        
        return first_name, last_name, ethnicity, gender
    
    def _get_fallback_first_name(self, gender: str, ethnicity: str) -> str:
        """Provide culturally appropriate fallback first names"""
        fallback_names = {
            'White': {
                'Male': ['James', 'Robert', 'John', 'Michael', 'William', 'David', 'Richard', 'Joseph'],
                'Female': ['Mary', 'Patricia', 'Jennifer', 'Linda', 'Elizabeth', 'Barbara', 'Susan', 'Jessica']
            },
            'Hispanic': {
                'Male': ['José', 'Luis', 'Miguel', 'Juan', 'Francisco', 'Alejandro', 'Antonio', 'Carlos'],
                'Female': ['María', 'Guadalupe', 'Juana', 'Margarita', 'Francisca', 'Teresa', 'Rosa', 'Ana']
            },
            'Black': {
                'Male': ['James', 'Robert', 'John', 'Michael', 'William', 'David', 'Richard', 'Charles'],
                'Female': ['Mary', 'Patricia', 'Linda', 'Elizabeth', 'Barbara', 'Susan', 'Jessica', 'Sarah']
            },
            'Asian': {
                'Male': ['Wei', 'Ming', 'Chen', 'Li', 'Kumar', 'Raj', 'Hiroshi', 'Takeshi'],
                'Female': ['Mei', 'Lin', 'Ling', 'Priya', 'Yuki', 'Soo', 'Hana', 'Anh']
            }
        }
        
        ethnicity_names = fallback_names.get(ethnicity, fallback_names['White'])
        return random.choice(ethnicity_names.get(gender, ethnicity_names['Male']))
    
    def _get_fallback_last_name(self, ethnicity: str) -> str:
        """Provide culturally appropriate fallback last names"""
        fallback_surnames = {
            'White': ['Smith', 'Johnson', 'Williams', 'Brown', 'Jones', 'Miller', 'Wilson', 'Moore'],
            'Hispanic': ['García', 'Rodríguez', 'Martínez', 'López', 'González', 'Pérez', 'Sánchez', 'Ramírez'],
            'Black': ['Johnson', 'Williams', 'Brown', 'Jones', 'Davis', 'Miller', 'Wilson', 'Moore'],
            'Asian': ['Li', 'Wang', 'Zhang', 'Liu', 'Chen', 'Yang', 'Huang', 'Zhao']
        }
        
        return random.choice(fallback_surnames.get(ethnicity, fallback_surnames['White']))
    
    def generate_random_name(self, gender: str = 'random') -> Tuple[str, str, str, str]:
        """Generate a random name with ethnicity based on US demographics"""
        ethnicity = self.get_weighted_ethnicity()
        return self.get_name_by_ethnicity(ethnicity, gender)
    
    def generate_multiple_names(self, count: int = 10, gender: str = 'random') -> List[Dict]:
        """Generate multiple random names"""
        names = []
        for _ in range(count):
            first, last, ethnicity, actual_gender = self.generate_random_name(gender)
            names.append({
                'full_name': f"{first} {last}",
                'first_name': first,
                'last_name': last,
                'ethnicity': ethnicity,
                'gender': actual_gender
            })
        return names
    
class GeneratePersona(Agent):
    agent_type: Literal["generate_persona"] = "generate_persona"  # type: ignore
    _agent_type: ClassVar[str] = agent_type  # Add this for registration
    agent_persist: bool = False

    system_prompt: str = generate_persona_system_prompt
    age: int = 0
    gender: str = ""
    username: str = ""
    first_name: str = ""
    last_name: str = ""
    full_name: str = ""
    ethnicity: str = ""

    generator: Any = Field(default=EthnicNameGenerator(), exclude=True)
    llm: Any = Field(default=None, exclude=True)
    model: str = Field(default=None, exclude=True)

    def randomize(self):
        self.age = random.randint(22, 67)
        # Use random.choices with explicit type casting to satisfy Literal type
        self.first_name, self.last_name, self.ethnicity, self.gender = self.generator.generate_random_name()
        self.full_name = f"{self.first_name} {self.last_name}"

    async def generate(
        self, llm: Any, model: str, 
        session_id: str, prompt: str, 
        tunables: Optional[Tunables] = None, 
        temperature=0.7
    ) -> AsyncGenerator[ChatMessage | ChatMessageStatus | ChatMessageError | ChatMessageStreaming, None]:
        self.randomize()

        original_prompt = prompt.strip()

        prompt = f"""\
```json
{json.dumps({
"age": self.age,
"gender": self.gender,
"ethnicity": self.ethnicity,
"full_name": self.full_name,
"first_name": self.first_name,
"last_name": self.last_name,
})}
```
"""
            
        if original_prompt:
            prompt += f"""
Incorporate the following into the job description: {original_prompt}
"""

        #
        # Generate the persona
        #
        logger.info(f"🤖 Generating persona for {self.full_name}")
        generating_message = None
        async for generating_message in self.llm_one_shot(
            llm=llm, model=model,
            session_id=session_id,
            prompt=prompt,
            system_prompt=generate_persona_system_prompt,
            temperature=temperature,
        ):
            if generating_message.status == ApiStatusType.ERROR:
                logger.error(f"Error generating persona: {generating_message.content}")
                yield generating_message
            if generating_message.status != ApiStatusType.DONE:
                yield generating_message
        
        if not generating_message:
            error_message = ChatMessageError(
                session_id=session_id,
                content="Persona generation failed to generate a response."
            )
            yield error_message
            return

        json_str = self.extract_json_from_text(generating_message.content)
        try:
            persona = {
                "age": self.age,
                "ethnicity": self.ethnicity,
                "gender": self.gender,
                "full_name": self.full_name,
                "first_name": self.first_name,
                "last_name": self.last_name,
            } | json.loads(json_str)
            self.username = persona.get("username", None)
            if not self.username:
                raise ValueError("LLM did not generate a username")
            self.username = re.sub(r'\s+', '.', self.username)
            user_dir = os.path.join(defines.user_dir, persona["username"])
            while os.path.exists(user_dir):
                match = re.match(r"^(.*?)(\d*)$", persona["username"])
                if match:
                    base = match.group(1)
                    num = match.group(2)
                    iteration = int(num) + 1 if num else 1
                    persona["username"] = f"{base}{iteration}"
                    user_dir = os.path.join(defines.user_dir, persona["username"])

            for key in persona:
                if isinstance(persona[key], str):
                    persona[key] = persona[key].strip()
            # Mark this persona as AI generated
            for to_lower in ["gender", "ethnicity"]:
                if to_lower in persona:
                    persona[to_lower] = persona[to_lower].lower()
            # Convert CITY, STATE, COUNTRY to Location type
            if "location" in persona:
                location_parts = persona["location"].split(",")
                if len(location_parts) == 3:
                    city, state, country = [part.strip() for part in location_parts]
                    persona["location"] = {
                        "city": city,
                        "state": state,
                        "country": country
                    }
                else:
                    logger.error(f"Invalid location format: {persona['location']}")
                    persona["location"] = None
            persona["is_ai"] = True
        except Exception as e:
            error_message = ChatMessageError(
                session_id=session_id,
                content=f"Error parsing LLM response: {str(e)}\n\n{json_str}"
            )
            logger.error(f"❌ Error parsing LLM response: {error_message.content}")
            logger.error(traceback.format_exc())
            logger.error(generating_message.content)
            yield error_message
            return

        logger.info(f"✅ Persona for {persona['username']} generated successfully")

        # Persona generated
        persona_message = ChatMessage(
            session_id=session_id,
            status=ApiStatusType.DONE,
            type=ApiMessageType.JSON,
            content = json.dumps(persona)
        )
        yield persona_message

        #
        # Generate the resume
        #
        status_message = ChatMessageStatus(
            session_id=session_id,
            activity = ApiActivityType.THINKING,
            content = f"Generating resume for {persona['full_name']}..."
        )
        logger.info(f"🤖 {status_message.content}")
        yield status_message

        content = f"""
```json
{{
    "full_name": "{persona["full_name"]}",
    "location": "{persona["location"]}",
    "age": {persona["age"]},
    "description": {persona["description"]},
    "title": {persona["title"]},
    "email": {persona["email"]},
    "phone": {persona["phone"]}
}}
```
"""
        if original_prompt:
            content += f"""
Make sure at least one of the candidate's job descriptions take into account the following: {original_prompt}."""

        async for generating_message in self.llm_one_shot(
            llm=llm, model=model,
            session_id=session_id,
            prompt=content,
            system_prompt=generate_resume_system_prompt,
            temperature=temperature,
        ):
            if generating_message.status == ApiStatusType.ERROR:
                logger.error(f"❌ Error generating resume: {generating_message.content}")
                raise Exception(generating_message.content)
            if generating_message.status != ApiStatusType.DONE:
                yield generating_message
        
        if not generating_message:
            error_message = ChatMessageError(
                session_id=session_id,
                content="Resume generation failed to generate a response."
            )
            logger.error(f"❌ {error_message.content}")
            yield error_message
            return

        resume = self.extract_markdown_from_text(generating_message.content)
        resume_message = ChatMessage(
            session_id=session_id,
            status=ApiStatusType.DONE,
            type=ApiMessageType.TEXT,
            content=resume
        )
        yield resume_message
        return

    def extract_json_from_text(self, text: str) -> str:
        """Extract JSON string from text that may contain other content."""
        json_pattern = r"```json\s*([\s\S]*?)\s*```"
        match = re.search(json_pattern, text)
        if match:
            return match.group(1).strip()

        # Try to find JSON without the markdown code block
        json_pattern = r"({[\s\S]*})"
        match = re.search(json_pattern, text)
        if match:
            return match.group(1).strip()

        raise ValueError("No JSON found in the response")

    def extract_markdown_from_text(self, text: str) -> str:
        """Extract Markdown string from text that may contain other content."""
        markdown_pattern = r"```(md|markdown)\s*([\s\S]*?)\s*```"
        match = re.search(markdown_pattern, text)
        if match:
            return match.group(2).strip()

        raise ValueError("No Markdown found in the response")
    
# Register the base agent
agent_registry.register(GeneratePersona._agent_type, GeneratePersona)
