import streamlit as st
import os
from groq import Groq
from typing import List, Dict, Optional, Union
import json
from datetime import datetime
import time
from functools import lru_cache
import logging
from contextlib import contextmanager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
MAX_RETRIES = 3
RETRY_DELAY = 1
DEFAULT_TEMPERATURE = 0.5
MAX_TOKENS = 1024
MODEL_NAME = "llama3-8b-8192"

class APIError(Exception):
    """Custom exception for API-related errors"""
    pass

class JSONParsingError(Exception):
    """Custom exception for JSON parsing errors"""
    pass

@contextmanager
def error_handler(context: str):
    """Context manager for handling errors with specific context"""
    try:
        yield
    except Exception as e:
        logger.error(f"Error in {context}: {str(e)}")
        st.error(f"An error occurred in {context}. Please try again.")
        raise

@st.cache_resource
def get_groq_client() -> Groq:
    """Initialize and cache Groq client"""
    try:
        return Groq(api_key=st.secrets["groq_api_key"])
    except Exception as e:
        logger.error(f"Failed to initialize Groq client: {str(e)}")
        st.error("Failed to initialize AI service. Please check your API key.")
        raise APIError("Failed to initialize Groq client")

class ContentAnalysisAgent:
    def __init__(self):
        """Initialize the agent with Groq client and default settings"""
        self.client = get_groq_client()
        self.system_prompt = """You are an expert social media content analyzer with deep understanding of engagement, 
        audience psychology, and content optimization. You must ALWAYS return responses in valid JSON format when requested.
        Analyze content step by step using a systematic approach."""
        
    @staticmethod
    def _display_thinking(thought: str):
        """Display agent's thinking process in a collapsible container"""
        with st.expander("🤔 Analysis Process", expanded=False):
            st.markdown(f"```\n{thought}\n```")

    def _call_api(self, messages: List[Dict], retries: int = MAX_RETRIES) -> Optional[str]:
        """Make API call with retry logic"""
        for attempt in range(retries):
            try:
                response = self.client.chat.completions.create(
                    messages=messages,
                    model=MODEL_NAME,
                    temperature=DEFAULT_TEMPERATURE,
                    max_tokens=MAX_TOKENS
                )
                return response.choices[0].message.content
            except Exception as e:
                if attempt == retries - 1:
                    logger.error(f"API call failed after {retries} attempts: {str(e)}")
                    raise APIError(f"Failed to get response from AI service: {str(e)}")
                time.sleep(RETRY_DELAY)
        return None

    @staticmethod
    def _parse_json(response: str) -> Dict:
        """Parse JSON from response with enhanced error handling"""
        try:
            # First attempt: direct JSON parsing
            return json.loads(response)
        except json.JSONDecodeError:
            try:
                # Second attempt: extract JSON structure
                start_idx = response.find('{')
                end_idx = response.rfind('}') + 1
                if start_idx != -1 and end_idx > start_idx:
                    json_str = response[start_idx:end_idx]
                    # Clean up common formatting issues
                    json_str = (json_str.replace('\n', ' ')
                              .replace('```json', '')
                              .replace('```', '')
                              .strip())
                    return json.loads(json_str)
            except (json.JSONDecodeError, ValueError) as e:
                logger.warning(f"JSON parsing failed: {str(e)}")
                # Return fallback structure
                return {
                    "style": "unknown",
                    "tones": ["neutral"],
                    "rating": "3",
                    "engagement_score": "50",
                    "analysis": {
                        "strengths": ["Content provided"],
                        "improvements": ["Format needs review"],
                        "audience_fit": "medium"
                    },
                    "error": "Response parsing failed"
                }

    def analyze_post(self, post_text: str) -> Dict:
        """Analyze post content with comprehensive error handling"""
        analysis_prompt = f"""Analyze this social media post and return ONLY a valid JSON object:
POST: {post_text}

Required structure:
{{
    "style": "posting style",
    "tones": ["tone1", "tone2"],
    "rating": "1-5",
    "engagement_score": "0-100",
    "analysis": {{
        "strengths": ["strength1", "strength2"],
        "improvements": ["improvement1", "improvement2"],
        "audience_fit": "low/medium/high"
    }}
}}"""

        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": analysis_prompt}
        ]

        with st.spinner("🔍 Analyzing content..."):
            try:
                analysis_response = self._call_api(messages)
                if not analysis_response:
                    raise APIError("No response received from API")
                
                analysis_result = self._parse_json(analysis_response)
                
                # Get recommendations
                recommendation_prompt = """Provide exactly 3 specific, actionable recommendations 
                to improve engagement. Return as a JSON array of strings."""
                
                messages.append({"role": "user", "content": recommendation_prompt})
                recommendations = self._call_api(messages)
                
                if recommendations:
                    try:
                        parsed_recommendations = json.loads(recommendations)
                        if isinstance(parsed_recommendations, list):
                            analysis_result["recommendations"] = parsed_recommendations
                        else:
                            analysis_result["recommendations"] = [recommendations.strip()]
                    except json.JSONDecodeError:
                        analysis_result["recommendations"] = [recommendations.strip()]
                
                return analysis_result
                
            except Exception as e:
                logger.error(f"Analysis failed: {str(e)}")
                st.error("Analysis failed. Please try again.")
                return None

class GraicieApp:
    def __init__(self):
        """Initialize the Graicie application"""
        self.agent = ContentAnalysisAgent()
        self.example_posts = {
            "Viral Marketing": "🚀 HUGE ANNOUNCEMENT! After months of work, my online course is finally LIVE! 🎉\n"
                             "Learn how I grew from 0 to 100K followers in 6 months! Early bird pricing ends tomorrow! 💫\n"
                             "#socialmedia #digitalmarketing #success",
            "Personal Story": "Sometimes life throws you curveballs... Today I faced my biggest fear and went "
                            "skydiving! 🪂 Swipe to see my reaction! Remember: growth happens outside your comfort zone 💕\n"
                            "#personalgrowth #motivation",
            "Educational": "🧠 5 Python Tips You Didn't Know:\n1. List comprehensions\n2. f-strings\n3. Walrus operator\n"
                         "4. Context managers\n5. Lambda functions\nSave this for later! 💡\n#coding #programming"
        }

    def _display_header(self):
        """Display application header"""
        st.title("🤖 Project Graicie - Advanced Content Analyzer")
        st.markdown("""
        ### Powered by LLaMA 3 & Agentic AI
        Get deep, AI-powered insights into your social media content using advanced language models.
        """)

    def _display_metrics(self, results: Dict):
        """Display analysis metrics in a structured format"""
        if not results:
            return

        # Main metrics
        cols = st.columns(4)
        with cols[0]:
            st.metric("Style", results["style"])
        with cols[1]:
            st.metric("Engagement", f"{results['engagement_score']}/100")
        with cols[2]:
            st.metric("Rating", f"{results['rating']}/5")
        with cols[3]:
            st.metric("Audience Fit", results["analysis"]["audience_fit"])

        # Content tones
        st.subheader("📊 Content Tones")
        tone_html = " ".join([
            f"<span style='background-color: black ; padding: 4px 8px; "
            f"margin: 4px; border-radius: 12px;'>{tone}</span>"
            for tone in results["tones"]
        ])
        st.markdown(tone_html, unsafe_allow_html=True)

        # Analysis details
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("💪 Strengths")
            for strength in results["analysis"]["strengths"]:
                st.markdown(f"✅ {strength}")

        with col2:
            st.subheader("🎯 Areas to Improve")
            for improvement in results["analysis"]["improvements"]:
                st.markdown(f"📌 {improvement}")

        # Recommendations
        if "recommendations" in results:
            st.subheader("🚀 Specific Recommendations")
            for idx, rec in enumerate(results["recommendations"], 1):
                st.markdown(f"{idx}. {rec}")

    def _display_sidebar(self):
        """Display sidebar with tips and information"""
        with st.sidebar:
            st.subheader("💡 Pro Tips")
            st.info("""
            **Content Best Practices:**
            1. Tell authentic stories
            2. Use relevant hashtags
            3. Include call-to-actions
            4. Add visual elements
            5. Engage with questions
            """)

            st.markdown("### 📊 Optimal Post Elements")
            st.markdown("""
            - Length: 80-150 characters
            - Hashtags: 3-5 relevant tags
            - Emojis: 2-3 key emojis
            - CTA: One clear action
            """)

    def run(self):
        """Run the Graicie application"""
        self._display_header()
        self._display_sidebar()

        # Main content area
        col1, col2 = st.columns([2, 1])

        with col1:
            st.subheader("📱 Try an Example Post")
            selected_example = st.selectbox(
                "Select an example:", 
                list(self.example_posts.keys())
            )

            if selected_example:
                example_text = self.example_posts[selected_example]
                st.text_area("Example Post", example_text, height=100, disabled=True)
                if st.button("Analyze Example", use_container_width=True):
                    with error_handler("example analysis"):
                        results = self.agent.analyze_post(example_text)
                        self._display_metrics(results)

            st.subheader("📝 Analyze Your Post")
            user_post = st.text_area(
                "Enter your post content:",
                height=150,
                placeholder="Type or paste your content here..."
            )

            if st.button("🔍 Analyze My Post", use_container_width=True):
                if user_post:
                    with error_handler("user post analysis"):
                        results = self.agent.analyze_post(user_post)
                        self._display_metrics(results)
                else:
                    st.warning("Please enter some content to analyze!")

        # Footer
        st.markdown(
            """
            <div style='text-align: center; padding: 20px;'>
                <p style='color: #666;'>
                    Powered by LLaMA 3 & Groq | Made with ❤️ by Project Graicie Team | 
                    © 2024 Project Graicie
                </p>
            </div>
            """,
            unsafe_allow_html=True,
        )

if __name__ == "__main__":
    try:
        app = GraicieApp()
        app.run()
    except Exception as e:
        logger.error(f"Application failed to start: {str(e)}")
        st.error("Application failed to start. Please check the logs.")