from typing import Type
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
import dspy
import os
from dotenv import load_dotenv


load_dotenv(dotenv_path="/home/azureuser/microlearn/backend/prompt_opt_dspy/.env")

azure_api_key = os.getenv("AZURE_API_KEY")
azure_api_host = os.getenv("AZURE_API_HOST")
azure_deployment_id = os.getenv("AZURE_DEPLOYMENT_ID")
azure_api_version = os.getenv("AZURE_API_VERSION")

# === Configure Azure OpenAI LM ===
lm = dspy.LM(
    "azure/Csqr-gpt-4o-mini",
    api_key=azure_api_key,
    api_base=azure_api_host,
    api_version=azure_api_version,
    temperature=0.0,
)
dspy.configure(lm=lm)


# === DSPy Signature ===
class RestructureSignature(dspy.Signature):
    """
    Restructure like a syllabus for the user journey or update the existing user journeys based on feedback.
    read the content {content} and make a new user journey if the feedack is "First iteration", Otherwise update the user journey in the {content} according to the user feedback {feedback}
    The user's company context is as follows:
        Company Name : {Company_Name}
        Industry : {Industry}
        Company Size : {Company_Size}
        Business Model : {Business_Model}
        Top Use Cases : {Top_Use_Cases}
        Learning Focus : {Learning_Focus}
        Tech Stack : {Tech_stack}
        Compliance : {Compliance}
    
    # User journey Creation Task
    
    ## Objective
    You are an intelligent and experienced data restructuring agent.
    Your task is to review the researched content provided and make a structured user journey.
    
    You must strictly adhere to the following guidelines.
    ## Guidlines
    - Use the level {Level} and user's experience with the topic {Experience} to decide the depth and balance of the user journey.
    - The users motive for taking this micro learning user journey is {motive}.
    - The user's knowledge level to different skills are as follows: (In percentage)
      1. Knowlege in SQL - {SQL}
      2. Data visualisation skill - {Viz}
      3. Statistical Anlysis skill - {Stats}
      4. Business Communication skill - {Comm}
      5. Data storytelling skill - {Story}
    - Set the content accordingly.
    - Divide the provided content to different stages
    - The toatal available time to cover the user journey is {available_time}.
    - **The sum of all stage durations MUST equal {available_time}. Do not exceed this.**
    - Include case studies where relevant, but **no projects**.
    - Based on the total available time, logically organize the topics into a structured, multi-stage user journey.
    - Distribute the available time ({available_time}) across the stages depending on content, focus and complexity.
    - Ensure the **sum of durations** across **all stages** is **{available_time}**
    - The user is either non-technical or minimally technical.
    
    ## Additional Instructions
    - **Before generating the stages, calculate the total duration and verify that the sum equals {available_time}.**
    - If necessary, adjust durations proportionally to maintain balance and alignment with the available time.
    - Do not comment or explain anything outside of the user journey content.
 
    ## Here is the researched content :
     {content}

    
    A detailed, structured user journey for a user on the given content, tailored to the {Level} level.
    The user is either a non technical or minimally technical person.
    - Generate a compelling main heading based on the user's Job Title and Topic.
      - The heading should reflect the relevance of the topic in the context of the user's role and the user's industry.
      - Keep it professional, engaging, and suitable.
    Job Title: {Job_Title}
    

    Divide the user journey into multiple stages.
    Each stage should contain:
      1. **heading** for the stage (include the stage number)
      2. **Focus**
      3. **Outcome**
      4. **Duration**
      5. **Topics Covered** — select relevant topics from the researched user journey content.
    The final stage should focus on application of the user's existing skills.
    Ensure the sum of durations across all stages is {available_time}.
    Do not include any project in the user journey, can include case studies.
    Should not comment anything other than the user journey.
    Format as markdown without '```'.
    Ensure the sum of durations across all stages is **exactly {available_time}**.
 
    Example output with two stages:
    # **Main Heading**
    ## **Stage1-Sub Heading**
    ### **Focus**
      Focus of stage1
    ### **Outcome**
      Expected outcome of stage 1
    ### **Duration**
      Duration1(Specify **ONLY** the time required to complete the stage1)
    ### **Topics Covered**
      - Topic1
      - Topic2
    ## **Stage2-Sub Heading**
    ### **Focus**
      Focus of stage2
    ### **Outcome**
      Expected outcome of stage 2
    ### **Duration**
      Duration2(Specify **ONLY** the time required to complete the stage2)
    ### **Topics Covered**
      - Topic1
      - Topic2
    Note:
    Duration1 + Duration2 should be equal to {available_time}
    Example output with three stages:
    # **Main Heading**
    ## **Stage1-Sub Heading**
    ### **Focus**
      Focus of stage1
    ### **Outcome**
      Expected outcome of stage 1
    ### **Duration**
      Duration1(Specify **ONLY** the time required to complete the stage1)
    ### **Topics Covered**
      - Topic1
      - Topic2
    ## **Stage2-Sub Heading**
    ### **Focus**
      Focus of stage2
    ### **Outcome**
      Expected outcome of stage 2
    ### **Duration**
      Duration2(Specify **ONLY** the time required to complete the stage2)
    ### **Topics Covered**
      - Topic1
      - Topic2
    ## **Stage3-Sub Heading**
    ### **Focus**
      Focus of stage3
    ### **Outcome**
      Expected outcome of stage 3
    ### **Duration**
      Duration3(Specify **ONLY** the time required to complete the stage2)
    ### **Topics Covered**
      - Topic1
      - Topic2
    Note:
    Duration1 + Duration2 + Duration3 should be equal to {available_time}
    ## Important:
    - **Sum of durations must equal {available_time}.**
    - No projects, but case studies allowed.
    - Markdown format only, no extra text or commentary.     

    """
    content: str = dspy.InputField()
    feedback: str = dspy.InputField()
    company_context: str = dspy.InputField()
    user_context: str = dspy.InputField()
    restructured_journey: str = dspy.OutputField(
        desc="Structured multi-stage user journey in markdown format with durations summing exactly to available_time"
    )

# === DSPy Module ===
class RestructureModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.restructure_predictor = dspy.Predict(RestructureSignature)
        predictor = dspy.Predict(RestructureSignature)
# --- Two Example Training Shots ---
        trainset = [
            dspy.Example(
                content='''
        - **Topic Area**: Fundamentals  
        - **Subtopics/Concepts**:  
            - Introduction to Artificial Intelligence in Healthcare  
            - Overview of Machine Learning and its Role in Patient Care  
            - Types of AI: Supervised, Unsupervised, and Reinforcement Learning  
            - Ethical Considerations in AI Applications  
            - Understanding Data Privacy: HIPAA and GDPR Compliance  

        - **Topic Area**: Applications  
        - **Subtopics/Concepts**:  
            - AI in Predictive Analytics: Enhancing Patient Outcomes  
            - AI in Patient Monitoring: Wearables and Remote Monitoring  
            - Use of Natural Language Processing (NLP) in Clinical Documentation  
            - Chatbots and Virtual Health Assistants for Patient Engagement  
            - Remote Diagnosis and Telemedicine Solutions  

        - **Topic Area**: Tools and Techniques  
        - **Subtopics/Concepts**:  
            - Overview of Data Analytics Tools: Python and PyTorch for healthcare  
            - Using AWS HealthLake for Data Management in Healthcare  
            - Data Visualization Tools for Presenting AI Insights  
            - Techniques in Data Preprocessing and Feature Engineering  
            - Implementing AI Models for Real-Time Patient Monitoring  

        - **Topic Area**: Advanced Topics  
        - **Subtopics/Concepts**:  
            - Deep Learning Applications in Medical Imaging  
            - AI-Driven Decision Support Systems for Healthcare Professionals  
            - Integrating AI into Existing Clinical Workflows  
            - Case Studies on AI Success Stories in Patient Care  
            - Challenges and Limitations of AI in Healthcare Practice  

        - **Topic Area**: Current Trends  
        - **Subtopics/Concepts**:  
            - Rise of Telehealth Solutions and AI Integration  
            - Innovations in Wearable Health Technology and AI  
            - The Role of AI in Addressing Pandemic Challenges  
            - The Future of AI in Personalized Medicine  
            - Collaborations Between Tech Companies and Healthcare Institutions  

        This structured user journey content is designed to provide a comprehensive overview of AI in Healthcare, particularly aligned with the interests and learning objectives of a Healthcare Analyst looking to explore how AI can improve patient care.
            
                ''',
                feedback="First iteration",
                company_context="""
                Company Name: MediCore Solutions
                Industry: Healthcare
                Company Size: 2000+ employees
                Business Model: B2B and B2C SaaS
                Top Use Cases: Predictive analytics, Patient monitoring
                Learning Focus: AI diagnostics and workflow automation
                Tech Stack: Python, PyTorch, AWS HealthLake
                Compliance: HIPAA, GDPR
                """,
                user_context="""
                Job Title: Healthcare Analyst
                Level: Intermediate
                Experience: 2 year
                motive: Exploring AI applications to improve patient care
                Knowledge in SQL: 70%
                Data Visualization Skill: 60%
                Statistical Analysis Skill: 40%
                Business Communication Skill: 50%
                Data Storytelling Skill: 45%
                Available Time: 240 minutes
                """,
                restructured_journey="""
                                
                # **Enhancing Patient Care through AI: A Healthcare Analyst’s Guide** 

                ## **Stage 1 - Introduction to AI in Healthcare**  
                ### **Focus**  
                Understanding the fundamental concepts of AI and its role in the healthcare industry.  
                ### **Outcome**  
                Establish a foundational knowledge of AI and machine learning principles, and comprehend ethical considerations and compliance regulations in healthcare.  
                ### **Duration**  
                60 minutes  
                ### **Topics Covered**  
                - Introduction to Artificial Intelligence in Healthcare  
                - Overview of Machine Learning and its Role in Patient Care  
                - Ethical Considerations in AI Applications  
                - Understanding Data Privacy: HIPAA and GDPR Compliance  

                ## **Stage 2 - AI Applications in Patient Care**  
                ### **Focus**  
                Exploring how AI is applied to diverse aspects of patient care, enhancing outcomes and engagement.  
                ### **Outcome**  
                Gain insights into practical applications of AI in predictive analytics, patient monitoring, and telemedicine, including the role of NLP and chatbots.  
                ### **Duration**  
                90 minutes  
                ### **Topics Covered**  
                - AI in Predictive Analytics: Enhancing Patient Outcomes  
                - AI in Patient Monitoring: Wearables and Remote Monitoring  
                - Use of Natural Language Processing (NLP) in Clinical Documentation  
                - Chatbots and Virtual Health Assistants for Patient Engagement  
                - Remote Diagnosis and Telemedicine Solutions  

                ## **Stage 3 - Tools, Techniques, and Advanced Topics**  
                ### **Focus**  
                Familiarizing with the technical tools and advanced applications of AI in healthcare while reviewing success stories.  
                ### **Outcome**  
                Understand data analytics tools, case studies of successful AI applications, and integration strategies for AI within existing clinical workflows.  
                ### **Duration**  
                90 minutes  
                ### **Topics Covered**  
                - Overview of Data Analytics Tools: Python and PyTorch for healthcare  
                - Using AWS HealthLake for Data Management in Healthcare  
                - Deep Learning Applications in Medical Imaging  
                - AI-Driven Decision Support Systems for Healthcare Professionals  
                - Case Studies on AI Success Stories in Patient Care  

                Total Duration: 60 + 90 + 90 = 240 minutes
                """
            ).with_inputs("content", "feedback", "company_context", "user_context"),

            dspy.Example(
                content="Research on SQL optimization techniques",
                feedback="Focus on advanced optimizations for experienced users.",
                company_context="""
                Company Name: FinNext Global
                Industry: FinTech
                Company Size: 1000-1500 employees
                Business Model: B2B
                Top Use Cases: Cross-border payments, Fraud prevention
                Learning Focus: Blockchain adoption and market opportunities
                Tech Stack: Ethereum, Hyperledger, Azure Blockchain Service
                Compliance: KYC, AML
                """,
                user_context="""
                Job Title: Business Development Executive
                Level: Beginner
                Experience: 5 years
                motive: Understanding blockchain for business strategy alignment
                Knowledge in SQL: 60%
                Data Visualization Skill: 70%
                Statistical Analysis Skill: 40%
                Business Communication Skill: 50%
                Data Storytelling Skill: 45%
                Available Time: 240 minutes
                """,
                restructured_journey="""
         # **Understanding Blockchain for Business Strategy Alignment in FinTech** 

        ## **Stage 1 - Introduction to Blockchain Fundamentals**  
        ### **Focus**  
        Introduce the fundamental concepts of blockchain technology.
        ### **Outcome**  
        Gain foundational knowledge about blockchain and its significance in the FinTech landscape.  
        ### **Duration**  
        60 minutes  
        ### **Topics Covered**  
        - Introduction to Blockchain  
        - Key Terms  
        - Differences between Blockchain Types  
        - Basic Mechanisms  

        ## **Stage 2 - Blockchain Applications in FinTech**  
        ### **Focus**  
        Explore specific blockchain applications relevant to FinTech, particularly in payments and fraud detection.  
        ### **Outcome**  
        Understand how blockchain enhances business processes within the FinTech sector.  
        ### **Duration**  
        60 minutes  
        ### **Topics Covered**  
        - Cross-border Payments  
        - Fraud Prevention  
        - Digital Identity Management  

        ## **Stage 3 - Tools and Platforms for Blockchain**  
        ### **Focus**  
        Familiarize with key blockchain platforms and tools that can be leveraged in business strategies.  
        ### **Outcome**  
        Identify relevant blockchain tools and platforms that can integrate into existing FinTech systems.  
        ### **Duration**  
        50 minutes  
        ### **Topics Covered**  
        - Overview of Ethereum and Hyperledger  
        - Azure Blockchain Service  
        - Integration with Existing Systems  

        ## **Stage 4 - Identifying Market Opportunities**  
        ### **Focus**  
        Analyze current trends, competition, and customer engagement opportunities within blockchain adoption.  
        ### **Outcome**  
        Develop insights on market opportunities and strategic advantages offered by blockchain.  
        ### **Duration**  
        40 minutes  
        ### **Topics Covered**  
        - Trends in Blockchain Adoption  
        - Competitive Landscape  
        - Customer Engagement  

        ## **Stage 5 - Business Strategy Alignment**  
        ### **Focus**  
        Align blockchain knowledge with business strategy and growth opportunities.  
        ### **Outcome**  
        Formulate strategic partnerships and business models leveraging blockchain technology.  
        ### **Duration**  
        30 minutes  
        ### **Topics Covered**  
        - Strategic Partnerships  
        - Business Models leveraging Blockchain  
        - Measuring ROI  

        ## **Stage 6 - Future Developments and Skills Application**  
        ### **Focus**  
        Understand future trends and how to continue developing skills in blockchain.  
        ### **Outcome**  
        Identify emerging trends and resources for continuous learning, and apply existing skills in statistical analysis and data visualization to business scenarios.  
        ### **Duration**  
        40 minutes  
        ### **Topics Covered**  
        - Emerging Trends  
        - Regulatory Landscape  
        - Education and Skill Development  

        **Total Duration: 240 minutes** 
                """
            ).with_inputs("content", "feedback", "company_context", "user_context"),
        ]

        # Metric: Check that duration is allocated in stages
        def restructure_metric(gold, pred, trace):
            return int("Stage" in pred.restructured_journey and str(gold.restructured_journey.count("Stage")) in pred.restructured_journey)
        # Compile with bootstrap few-shot optimizer
        optimizer = dspy.BootstrapFewShot(metric=restructure_metric)
        optimizer.compile(predictor, trainset=trainset)

        # Store optimized predictor
        self.restructure_predictor = predictor
        print("=== Optimized ===")
        print(predictor)

    def forward(self, content, feedback, company_context, user_context):
        return self.restructure_predictor(
            content=content,
            feedback=feedback,
            company_context=company_context,
            user_context=user_context
        )

# === Tool Input Schema ===
class RestructureToolInput(BaseModel):
    """Input schema for RestructureAnalysisTool."""
    content: str = Field(..., description="The researched content to structure into a user journey")
    feedback: str = Field(..., description="User feedback (First iteration or update request)")
    Company_Name: str = Field(..., description="Name of the company")
    Industry: str = Field(..., description="Industry of the company")
    Company_Size: str = Field(..., description="Size of the company")
    Business_Model: str = Field(..., description="Company's business model")
    Top_Use_Cases: str = Field(..., description="Top use cases the company focuses on")
    Learning_Focus: str = Field(..., description="Learning focus for the user")
    Tech_stack: str = Field(..., description="Company's tech stack")
    Compliance: str = Field(..., description="Relevant compliance requirements")
    Job_Title: str = Field(..., description="Job title of the user")
    Level: str = Field(..., description="Learning level (Beginner, Intermediate, Advanced)")
    Experience: str = Field(..., description="User's experience with the topic")
    motive: str = Field(..., description="User's motive for research")
    SQL: int = Field(..., description="Knowledge in SQL (percentage)")
    Viz: int = Field(..., description="Data visualization skill (percentage)")
    Stats: int = Field(..., description="Statistical analysis skill (percentage)")
    Comm: int = Field(..., description="Business communication skill (percentage)")
    Story: int = Field(..., description="Data storytelling skill (percentage)")
    available_time: int = Field(..., description="Total available time for the user journey (in minutes)")

# === Custom DSPy Tool ===
class RestructureAnalysisTool(BaseTool):
    name: str = "Restructure Analysis Tool"
    description: str = (
        "Restructures researched content into a structured user journey syllabus based on user feedback, "
        "company context, and user profile. Ensures durations sum exactly to available_time."
    )
    args_schema: Type[BaseModel] = RestructureToolInput

    def _run(
        self,
        content: str,
        feedback: str,
        Company_Name: str,
        Industry: str,
        Company_Size: str,
        Business_Model: str,
        Top_Use_Cases: str,
        Learning_Focus: str,
        Tech_stack: str,
        Compliance: str,
        Job_Title: str,
        Level: str,
        Experience: str,
        motive: str,
        SQL: int,
        Viz: int,
        Stats: int,
        Comm: int,
        Story: int,
        available_time: int
    ) -> str:

        # Build company + user context
        company_context = f"""
        Company Name: {Company_Name}
        Industry: {Industry}
        Company Size: {Company_Size}
        Business Model: {Business_Model}
        Top Use Cases: {Top_Use_Cases}
        Learning Focus: {Learning_Focus}
        Tech Stack: {Tech_stack}
        Compliance: {Compliance}
        """

        user_context = f"""
        Job Title: {Job_Title}
        Level: {Level}
        Experience: {Experience}
        Motive: {motive}
        Knowledge in SQL: {SQL}%
        Data Visualization Skill: {Viz}%
        Statistical Analysis Skill: {Stats}%
        Business Communication Skill: {Comm}%
        Data Storytelling Skill: {Story}%
        Available Time: {available_time}
        """

        restructure = RestructureModule()
        result = restructure.forward(content, feedback, company_context, user_context)

        return result.restructured_journey


# # from restructure_tool import RestructureAnalysisTool

# if __name__ == "__main__":
#     tool = RestructureAnalysisTool()

#     result = tool._run(
#         content="""

#     - **Topic Area**: Fundamentals of Cyber Security
#     - **Subtopics/Concepts**:
#         - Introduction to Cyber Security: Importance and relevance in EdTech
#         - Key Terminology: Threats, vulnerabilities, and risks
#         - Types of Cyber Attacks: Phishing, malware, ransomware
#         - Basic Security Principles: Confidentiality, integrity, availability (CIA triad)

#     - **Topic Area**: Cyber Security in Education Technology
#     - **Subtopics/Concepts**:
#         - Data Protection Regulations: Overview of GDPR and FERPA
#         - Best Practices for Data Security in EdTech: User data management and protection
#         - Case Studies: Successful implementation of security measures in EdTech companies

#     - **Topic Area**: Tools and Techniques for Cyber Security
#     - **Subtopics/Concepts**:
#         - Security Tools Overview: Firewalls, antivirus software, and intrusion detection systems
#         - Introduction to Encryption: Importance of data encryption in protecting sensitive information
#         - User Authentication Methods: Multi-factor authentication and secure password practices

#     - **Topic Area**: Current Trends in Cyber Security
#     - **Subtopics/Concepts**:
#         - AI and Cyber Security: How AI is transforming security measures in EdTech
#         - Emerging Threats: Understanding the evolving landscape of cyber threats
#         - Cyber Security Frameworks: Overview of NIST and ISO standards relevant to EdTech

#     - **Topic Area**: Upskilling for Cyber Security
#     - **Subtopics/Concepts**:
#         - Learning Resources: Online courses and certifications for beginners in cyber security
#         - Networking Opportunities: Joining EdTech and Cyber Security communities
#         - Building a Cyber Security Mindset: Importance of continuous learning and awareness
                
#                 """,
#         feedback="First iteration",
#         Company_Name="TechNova Inc",
#         Industry="EdTech",
#         Company_Size="500-1000 employees",
#         Business_Model="B2B SaaS",
#         Top_Use_Cases="Personalized learning, AI tutors",
#         Learning_Focus="AI adoption in education",
#         Tech_stack="Python, TensorFlow, Azure",
#         Compliance="GDPR, FERPA",
        
#         Job_Title="Product Manager",
#         Level="Beginner",
#         Experience="2 years working with SQL and dashboards",
#         motive="Upskilling for AI-driven product development",
#         SQL=70,
#         Viz=60,
#         Stats=40,
#         Comm=50,
#         Story=45,
#         available_time=240
#     )

#     print("=== Restructured User Journey ===")
#     print(result)
