import dspy
from typing import Optional, Type, TYPE_CHECKING, List, Any, TypeVar

if TYPE_CHECKING:
    from pydantic import BaseModel

# TypeVars for dynamic types
EntityT = TypeVar('EntityT')
RelationT = TypeVar('RelationT')
QueryT = TypeVar('QueryT')


class AI:
  def __init__(
    self,
    model: str = "gemini/gemini-2.5-flash",
    temperature: float = 0.0,
    api_key: str = None,
    api_base: str = None,
    entity_type: Type[EntityT] = None,
    relation_type: Type[RelationT] = None,
    query_type: Type[QueryT] = None,
    ontology: Any = None
  ):
    """Initialize with optional model configuration

    Args:
        model: Name of model to use (e.g. 'gpt-4')
        temperature: Temperature for model sampling
        api_key: API key for model access
        api_base: Specify the base URL endpoint for making API calls to a language model service
        entity_type: The Entity class type from KGMem
        relation_type: The Relation class type from KGMem
        query_type: The Query class type from KGMem
        ontology: The ontology object containing entity types and relation types
    """
    self.dspy = dspy
    self.model = model
    self.temperature = temperature
    self.api_key = api_key
    self.api_base = api_base
    self.entity_type = entity_type
    self.relation_type = relation_type
    self.query_type = query_type
    self.ontology = ontology
    self.init_model(model, temperature, api_key, api_base)

  
  def init_model(
    self,
    model: str = None,
    temperature: float = None,
    api_key: str = None,
    api_base: str = None
  ):
    """Initialize or reinitialize the model with new parameters

    Args:
        model: Name of model to use (e.g. 'gpt-4')
        temperature: Temperature for model sampling
        api_key: API key for model access
        api_base: API base for model access
    """
    # Update instance variables if new values provided
    if model is not None:
      self.model = model
    if temperature is not None:
      self.temperature = temperature
    if api_key is not None:
      self.api_key = api_key
    if api_base is not None:
      self.api_base = api_base

    # Initialize dspy LM with current settings
    if self.api_key:
      self.lm = dspy.LM(model=self.model, api_key=self.api_key, temperature=self.temperature, api_base=self.api_base, max_tokens=16384)
    else:
      self.lm = dspy.LM(model=self.model, temperature=self.temperature, api_base=self.api_base, max_tokens=16384)

    self.dspy.configure(lm=self.lm)

  def _validate_relation_type(self, obj: Any) -> bool:
    """Validate if an object is of the correct Relation type"""
    return self.relation_type and isinstance(obj, self.relation_type)

  def __call__(self, user_prompt: str, system_prompt: Optional[str] = None) -> str:
    class QA(dspy.Signature):
        __doc__ = system_prompt
        
        user_prompt: str = dspy.InputField(desc="The user's request.")
        response: str = dspy.OutputField(desc="The text response generated by the language model.")
        
    predict = dspy.Predict(QA)
    return predict(user_prompt=user_prompt).response
  
  def extract_relations(self, text_chunk: str) -> List[Any]:
    """Extract relations from a text chunk.
    
    Returns a list of Relation objects (the actual type is self.relation_type)
    """
    if not self.entity_type or not self.relation_type or not self.ontology:
      return []
    
    # Get entity and relation types from ontology
    entity_types = self.ontology.entity_types
    relation_types = self.ontology.relation_types
    
    if not entity_types or not relation_types:
      print("Warning: No entity types or relation types defined in ontology")
      return []
    
    # Define the extraction signature
    class ExtractKnowledge(dspy.Signature):
      """Extract entities and their relationships from text. Identify key entities and how they relate to each other."""
      
      text: str = dspy.InputField(desc="The text to extract knowledge from")
      entity_types_info: str = dspy.InputField(desc="Available entity types and their descriptions")
      relation_types_info: str = dspy.InputField(desc="Available relation types with their constraints") 
      extracted_data: str = dspy.OutputField(desc="JSON array of relations. Each relation has: entity0_name, entity0_type, entity1_name, entity1_type, relation_description, relation_type")
    
    # Build detailed type information including constraints
    entity_types_info = "Entity Types:\n" + "\n".join([f"- {et.name}" for et in entity_types])
    
    relation_types_info = "Relation Types:\n"
    for rt in relation_types:
      relation_types_info += f"- {rt.name}: connects {rt.entity_type0.name} to {rt.entity_type1.name}"
      if rt.context:
        relation_types_info += f" (context: {rt.context})"
      if rt.directed:
        relation_types_info += " [directed]"
      relation_types_info += "\n"
    
    # Use Chain of Thought for better extraction
    predict = dspy.ChainOfThought(ExtractKnowledge)
    result = predict(
      text=text_chunk,
      entity_types_info=entity_types_info,
      relation_types_info=relation_types_info
    )
    
    # Helper function to create or get entity from cache
    def get_or_create_entity(name: str, type_name: str, entity_cache: dict) -> Optional[Any]:
      """Get entity from cache or create new one if valid."""
      cache_key = (name, type_name)
      
      if cache_key in entity_cache:
        return entity_cache[cache_key]
      
      # Find the EntityType object
      entity_type_obj = next((et for et in entity_types if et.name == type_name), None)
      if not entity_type_obj:
        print(f"Warning: Entity type '{type_name}' not found in ontology")
        return None
      
      try:
        entity = self.entity_type(name=name, type=entity_type_obj)
        entity_cache[cache_key] = entity
        return entity
      except Exception as e:
        print(f"Error creating entity '{name}' of type '{type_name}': {e}")
        return None
    
    # Parse the extracted data
    relations = []
    entity_cache = {}
    
    try:
      import json
      extracted_items = json.loads(result.extracted_data)
      
      if not isinstance(extracted_items, list):
        print("Warning: Extracted data is not a list")
        return []
      
      for idx, item in enumerate(extracted_items):
        try:
          # Validate required fields
          required_fields = ['entity0_name', 'entity0_type', 'entity1_name', 
                           'entity1_type', 'relation_description', 'relation_type']
          missing_fields = [f for f in required_fields if f not in item]
          if missing_fields:
            print(f"Warning: Item {idx} missing fields: {missing_fields}")
            continue
          
          # Find the RelationType object
          relation_type_obj = next(
            (rt for rt in relation_types if rt.name == item['relation_type']), 
            None
          )
          if not relation_type_obj:
            print(f"Warning: Relation type '{item['relation_type']}' not found")
            continue
          
          # Validate entity types match relation type constraints
          if item['entity0_type'] != relation_type_obj.entity_type0.name:
            print(f"Warning: Entity0 type mismatch. Expected {relation_type_obj.entity_type0.name}, "
                  f"got {item['entity0_type']} for relation {item['relation_type']}")
            continue
            
          if item['entity1_type'] != relation_type_obj.entity_type1.name:
            print(f"Warning: Entity1 type mismatch. Expected {relation_type_obj.entity_type1.name}, "
                  f"got {item['entity1_type']} for relation {item['relation_type']}")
            continue
          
          # Get or create entities
          entity0 = get_or_create_entity(
            item['entity0_name'], 
            item['entity0_type'], 
            entity_cache
          )
          entity1 = get_or_create_entity(
            item['entity1_name'], 
            item['entity1_type'], 
            entity_cache
          )
          
          if not entity0 or not entity1:
            continue
          
          # Create the relation
          try:
            relation = self.relation_type(
              entity0=entity0,
              entity1=entity1,
              relation=item['relation_description'],
              type=relation_type_obj
            )
            relations.append(relation)
          except Exception as e:
            print(f"Error creating relation: {e}")
            
        except Exception as e:
          print(f"Error processing item {idx}: {e}")
          continue
          
    except json.JSONDecodeError as e:
      print(f"Failed to parse JSON from LLM response: {e}")
      print(f"Raw response: {result.extracted_data[:200]}...")  # Show first 200 chars
    except Exception as e:
      print(f"Unexpected error in extract_relations: {e}")
    
    print(f"Successfully extracted {len(relations)} relations from text")
    return relations
  
  def retrieve_relations(self, query_name: str, query_context: str, all_relations: List[Any], types_to_retrieve: list, limit: int = 10) -> List[Any]:
    """Retrieve relevant relations from the knowledge graph based on a query.
    
    Args:
        query_name: The name/description of the query
        query_context: The context in which this query is made
        all_relations: List of Relation objects (the actual type is self.relation_type)
        types_to_retrieve: List of entity/relation types to retrieve
        limit: Maximum number of relations to retrieve
    
    Returns:
        List of Relation objects (the actual type is self.relation_type)
    """
    
    class RetrieveRelations(dspy.Signature):
        """Given a query and a list of relations from a knowledge graph, identify and return the most relevant relations that match the query intent and type requirements."""
        
        query_name: str = dspy.InputField(desc="The name/description of the query")
        query_context: str = dspy.InputField(desc="The context in which this query is made")
        all_relations_str: str = dspy.InputField(desc="String representation of all available relations")
        types_to_retrieve_str: str = dspy.InputField(desc="String representation of entity/relation types to retrieve")
        limit: int = dspy.InputField(desc="Maximum number of relations to retrieve")
        relevant_indices: str = dspy.OutputField(desc="Comma-separated indices of the most relevant relations (e.g., '0,3,5,7')")
    
    # Convert relations and types to string representations for the LLM
    relations_str = "\n".join([f"{i}: {rel}" for i, rel in enumerate(all_relations)])
    types_str = ", ".join([str(t) for t in types_to_retrieve])
    
    predict = dspy.ChainOfThought(RetrieveRelations)
    result = predict(
        query_name=query_name,
        query_context=query_context,
        all_relations_str=relations_str,
        types_to_retrieve_str=types_str,
        limit=limit
    )
    
    # Parse the indices and return the corresponding relations
    try:
        indices = [int(idx.strip()) for idx in result.relevant_indices.split(',') if idx.strip()]
        # Ensure indices are valid
        indices = [idx for idx in indices if 0 <= idx < len(all_relations)]
        # Limit to requested number
        indices = indices[:limit]
        return [all_relations[idx] for idx in indices]
    except (ValueError, IndexError):
        # If parsing fails, return empty list
        return []
    