import asyncio
from typing import Dict, List
from langchain_community.chat_models import ChatOllama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage

class SimpleAnalysisAgent:
    """A simple agent that analyzes text and provides insights."""
    
    def __init__(self, model_name: str = "llama3.1:8b"):
        if not self.check_ollama_model(model_name):
            print(f"❌ Model {model_name} not found. Try: ollama pull {model_name}")
            raise ValueError(f"Model {model_name} not available")
        
        self.llm = ChatOllama(model=model_name, temperature=0.1)
        self.conversation_history = []
    
    def classify_text(self, text: str) -> str:
        """Classify the type of text."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "Classify this text as one of: news, blog, email, code, academic, or other. Respond with just the category."),
            ("human", "{text}")
        ])
        
        chain = prompt | self.llm
        result = chain.invoke({"text": text[:500]})  # Limit length
        return result.content.strip().lower()
    
    def extract_key_points(self, text: str) -> List[str]:
        """Extract key points from text."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "Extract 3-5 key points from this text. Return as a simple numbered list."),
            ("human", "{text}")
        ])
        
        chain = prompt | self.llm
        result = chain.invoke({"text": text})
        
        # Parse the numbered list
        lines = result.content.strip().split('\n')
        points = [line.strip() for line in lines if line.strip() and any(c.isdigit() for c in line[:3])]
        return points[:5]  # Limit to 5 points
    
    def summarize(self, text: str) -> str:
        """Create a summary of the text."""
        prompt = ChatPromptTemplate.from_messages([
            ("system", "Summarize this text in 2-3 sentences. Be concise and clear."),
            ("human", "{text}")
        ])
        
        chain = prompt | self.llm
        result = chain.invoke({"text": text})
        return result.content.strip()
    
    def analyze_text(self, text: str) -> Dict:
        """Complete analysis of text."""
        print(f"🔍 Analyzing text ({len(text)} characters)...")
        
        # Step 1: Classify
        print("  📊 Classifying...")
        category = self.classify_text(text)
        
        # Step 2: Extract key points
        print("  🔑 Extracting key points...")
        key_points = self.extract_key_points(text)
        
        # Step 3: Summarize
        print("  📝 Summarizing...")
        summary = self.summarize(text)
        
        return {
            "category": category,
            "key_points": key_points,
            "summary": summary,
            "length": len(text)
        }
    
    @staticmethod
    def check_ollama_model(model_name: str) -> bool:
        try:
            import requests
            response = requests.get("http://localhost:11434/api/tags")
            models = [m["name"] for m in response.json().get("models", [])]
            return model_name in models
        except:
            return False

def demo_agent():
    """Demonstrate the agent with sample texts."""
    print("🤖 LangChain Agent with Ollama Demo")
    print("=" * 40)
    
    try:
        agent = SimpleAnalysisAgent()
        print("✅ Agent initialized successfully\n")
    except ValueError as e:
        print(e)
        return
    
    # Sample texts to analyze
    samples = [
        {
            "name": "Tech News",
            "text": """
            Apple announced today that its new iPhone 15 will feature USB-C charging,
            marking a significant shift from the Lightning connector. The change comes
            after pressure from the European Union's new charging regulations. The new
            phones will also include improved cameras and faster processors. Industry
            analysts expect this to boost sales significantly in the next quarter.
            """
        },
        {
            "name": "Code Comment",
            "text": """
            def process_data(input_list):
                # This function takes a list of numbers and returns the average
                # It handles empty lists by returning 0
                if not input_list:
                    return 0
                return sum(input_list) / len(input_list)
            """
        }
    ]
    
    # Analyze each sample
    for sample in samples:
        print(f"📋 Sample: {sample['name']}")
        print("-" * 30)
        
        result = agent.analyze_text(sample["text"].strip())
        
        print(f"Category: {result['category']}")
        print(f"Summary: {result['summary']}")
        print("Key Points:")
        for point in result['key_points']:
            print(f"  • {point}")
        print()

demo_agent()
print("\n✅ Done!")

def interactive_mode():
    """Interactive mode for testing with your own text."""
    print("🔄 Interactive Mode")
    print("=" * 40)
    
    try:
        agent = SimpleAnalysisAgent()
    except ValueError as e:
        print(e)
        return
    
    print("Enter text to analyze (or 'quit' to exit):")
    
    while True:
        text = input("\n> ")
        
        if text.lower() in ['quit', 'exit', 'q']:
            break
        
        if len(text.strip()) < 10:
            print("Please enter more text (at least 10 characters)")
            continue
        
        result = agent.analyze_text(text)
        
        print(f"\n📊 Results:")
        print(f"Category: {result['category']}")
        print(f"Summary: {result['summary']}")
        if result['key_points']:
            print("Key Points:")
            for point in result['key_points']:
                print(f"  • {point}")

interactive_mode()