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

# === Load secrets ===
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_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 Feedback Updation Task
        
        ## Objective
        You are an intelligent and experienced restructuring agent.
        Your task is to review the already created user journey content {content}.
        Then, based on the user feedback — "{feedback}", update and restructure the user journey accordingly.
        
        ## Guidelines
        - Read and understand the original user journey content.
        - Understand the user feedback carefully.
        - If the feedback asks for less content, reduce topics proportionally.
        - If it asks for more depth, expand topics proportionally.
        - Do not adjust durations.
        
        ## Additional Instructions
        - Keep the format and markdown structure identical to the original.
        - No projects, but you can include case studies.
        - Return only the updated markdown, no additional explanation.

        ## Here is the already existing user journey :
        {content}
        
    The updated user journey markdown content as per feedback.

    """

    company_context: str = dspy.InputField()
    content: str = dspy.InputField()
    feedback: str = dspy.InputField()
    updated_content: str = dspy.OutputField(
        desc="The updated user journey markdown syllabus"
    )


# === DSPy Module ===
class RestructureModule(dspy.Module):
    def __init__(self):
        super().__init__()
        self.predictor = dspy.Predict(RestructureSignature)

        # --- Few-shot examples ---
        trainset = [
            dspy.Example(
                company_context="""
                    Company Name: FinSight, 
                    Industry: Finance, 
                    Company Size: 250, 
                    Business Model: B2B, 
                    Top Use Cases: Risk modeling, 
                    Learning Focus: Data Analysis, 
                    Tech Stack: SQL, Tableau, 
                    Compliance: SOX
                """,
                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="Add more depth on NLP",
                updated_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  
                        - Definition and Importance of NLP in Healthcare  
                        - Techniques in NLP: Tokenization, Named Entity Recognition, and Sentiment Analysis  
                        - Real-world Applications: Automating Clinical Note Generation and Patient Interaction  
                        - Case Studies: Successful Implementations of NLP in Healthcare Settings  
                    - 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  

                """
            ).with_inputs("company_context", "content", "feedback"),
            dspy.Example(
                company_context="""
                    Company Name: HealthAI, 
                    Industry: Healthcare, 
                    Company Size: 500, 
                    Business Model: SaaS, 
                    Top Use Cases: Patient analytics, 
                    Learning Focus: Machine Learning, 
                    Tech Stack: Python, AWS, 
                    Compliance: HIPAA
                """,
                content="""
                    
            - **Topic Area**: Fundamentals  
            - **Subtopics/Concepts**:  
                - Introduction to Blockchain: Understanding what blockchain technology is and its significance in the FinTech landscape.  
                - Key Terms: Decentralization, ledgers, nodes, smart contracts, tokens, hash functions.  
                - Differences between Blockchain Types: Public vs. Private vs. Consortium blockchains.  
                - Basic Mechanisms: How transactions are processed and validated in a blockchain network.  
                - Overview of Cryptography: Basic principles of how cryptography secures blockchain solutions.  

            - **Topic Area**: Applications in FinTech  
            - **Subtopics/Concepts**:  
                - Cross-border Payments: How blockchain facilitates faster, cheaper, and transparent international money transfers.  
                - Fraud Prevention: Use of blockchain to enhance security and traceability of transactions.  
                - Digital Identity Management: Evaluating blockchain's role in KYC and AML compliance processes for identity verification.  
                - Smart Contracts: Understanding automated contract execution and its applications in financial services. 

            - **Topic Area**: Tools & Platforms  
            - **Subtopics/Concepts**:  
                - Overview of Ethereum and Hyperledger: Key differences, use cases, and their relevance in FinTech.  
                - Azure Blockchain Service: Brief introduction to Microsoft’s blockchain service and its business applications.  
                - Wallet Solutions: Insight into digital wallets and their importance for businesses.  
                - Integration with Existing Systems: Exploring middleware and API solutions to connect blockchain with legacy systems.  

            - **Topic Area**: Market Opportunities  
            - **Subtopics/Concepts**:  
                - Trends in Blockchain Adoption: Current and future trends driving blockchain in FinTech, including regulatory factors.  
                - Competitive Landscape: Overview of companies and startups utilizing blockchain for competitive advantage.  
                - Customer Engagement: How blockchain can improve client engagement through transparency and trust.  

            - **Topic Area**: Business Strategy Alignment  
            - **Subtopics/Concepts**:  
                - Strategic Partnerships: Identifying potential partnerships in the blockchain ecosystem for business growth.  
                - Business Models leveraging Blockchain: Understanding various B2B models, including payment solutions, lending, and insurance.  
                - Measuring ROI: Exploring metrics and frameworks to assess the financial impact of blockchain initiatives.  

            - **Topic Area**: Future Developments  
            - **Subtopics/Concepts**:  
                - Emerging Trends: Overview of upcoming blockchain technology trends such as NFTs, CBDCs (Central Bank Digital Currencies), and DeFi (Decentralized Finance).  
                - Regulatory Landscape: Understanding evolving compliance standards and regulations impacting blockchain implementations.  
                - Education and Skill Development: Resources and communities to explore for continuous learning in blockchain and FinTech.

            This structured user journey provides the necessary foundation and insight for a Business Development Executive at FinNext Global to understand blockchain's application in developing business strategies and enhancing client engagement in the FinTech sector.

                
                """,
                feedback="First iteration",
                updated_content="""

            - **Topic Area**: Fundamentals  
            - **Subtopics/Concepts**:  
                - Introduction to Blockchain: Understanding what blockchain technology is and its significance in the FinTech landscape.  
                - Key Terms: Decentralization, ledgers, nodes, smart contracts, tokens, hash functions.  
                - Differences between Blockchain Types: Public vs. Private vs. Consortium blockchains.  
                - Basic Mechanisms: How transactions are processed and validated in a blockchain network.  
                - Overview of Cryptography: Basic principles of how cryptography secures blockchain solutions.  

            - **Topic Area**: Applications in FinTech  
            - **Subtopics/Concepts**:  
                - Cross-border Payments: How blockchain facilitates faster, cheaper, and transparent international money transfers.  
                - Fraud Prevention: Use of blockchain to enhance security and traceability of transactions.  
                - Digital Identity Management: Evaluating blockchain's role in KYC and AML compliance processes for identity verification.  
                - Smart Contracts: Understanding automated contract execution and its applications in financial services. 
                - **Natural Language Processing (NLP)**: Exploring how NLP can integrate with blockchain for enhanced transaction analysis and fraud detection.  
                
            - **Topic Area**: Tools & Platforms  
            - **Subtopics/Concepts**:  
                - Overview of Ethereum and Hyperledger: Key differences, use cases, and their relevance in FinTech.  
                - Azure Blockchain Service: Brief introduction to Microsoft’s blockchain service and its business applications.  
                - Wallet Solutions: Insight into digital wallets and their importance for businesses.  
                - Integration with Existing Systems: Exploring middleware and API solutions to connect blockchain with legacy systems.  
                - **NLP Tools in Blockchain**: Identifying tools that utilize NLP for analyzing blockchain transaction data and enhancing customer interactions.  

            - **Topic Area**: Market Opportunities  
            - **Subtopics/Concepts**:  
                - Trends in Blockchain Adoption: Current and future trends driving blockchain in FinTech, including regulatory factors.  
                - Competitive Landscape: Overview of companies and startups utilizing blockchain for competitive advantage.  
                - Customer Engagement: How blockchain can improve client engagement through transparency and trust.  
                - **Leveraging NLP for Market Analysis**: Analyzing how NLP can interpret market data and client sentiment for strategic business insights.  

            - **Topic Area**: Business Strategy Alignment  
            - **Subtopics/Concepts**:  
                - Strategic Partnerships: Identifying potential partnerships in the blockchain ecosystem for business growth.  
                - Business Models leveraging Blockchain: Understanding various B2B models, including payment solutions, lending, and insurance.  
                - Measuring ROI: Exploring metrics and frameworks to assess the financial impact of blockchain initiatives.  
                - **NLP in Business Development**: Using NLP techniques to analyze client feedback and enhance service offerings.  

            - **Topic Area**: Future Developments  
            - **Subtopics/Concepts**:  
                - Emerging Trends: Overview of upcoming blockchain technology trends such as NFTs, CBDCs (Central Bank Digital Currencies), and DeFi (Decentralized Finance).  
                - Regulatory Landscape: Understanding evolving compliance standards and regulations impacting blockchain implementations.  
                - Education and Skill Development: Resources and communities to explore for continuous learning in blockchain and FinTech.  
                - **Future of NLP in FinTech**: Discussing potential advancements in NLP and its applications in customer service and fraud detection within blockchain frameworks.

            This structured user journey provides the necessary foundation and insight for a Business Development Executive at FinNext Global to understand blockchain's application in developing business strategies, enhancing client engagement in the FinTech sector, and leveraging NLP for deeper market insights and operational efficiency.

                """
            ).with_inputs("company_context", "content", "feedback"),
        ]

        # --- Metric: ensure markdown structure with '## User Journey' ---
        def restructure_metric(gold, pred, trace=None):
            output = getattr(pred, "updated_content", "") or ""
            return int(output.strip().startswith("## User Journey"))

        optimizer = dspy.BootstrapFewShot(metric=restructure_metric)
        self.optimized_predictor = optimizer.compile(
            student=self.predictor,
            trainset=trainset
        )

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


# === Tool Input Schema ===
class RestructureInput(BaseModel):
    """Input schema for RestructureTool."""
    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")
    content: str = Field(..., description="Existing user journey content in markdown")
    feedback: str = Field(..., description="User feedback for updating the journey")


# === Custom DSPy Tool ===
class RestructureTool(BaseTool):
    name: str = "Restructure Tool"
    description: str = (
        "Restructures or updates a user journey syllabus based on content and user feedback."
    )
    args_schema: Type[BaseModel] = RestructureInput

    def _run(
        self,
        Company_Name: str,
        Industry: str,
        Company_Size: str,
        Business_Model: str,
        Top_Use_Cases: str,
        Learning_Focus: str,
        Tech_stack: str,
        Compliance: str,
        content: str,
        feedback: str
    ) -> str:

        # Build company context string
        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}
        """

        restructure_module = RestructureModule()
        result = restructure_module.forward(
            company_context=company_context,
            content=content,
            feedback=feedback
        )

        return result.updated_content




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

# # === Load secrets ===
# 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_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 Feedback Updation Task
        
#         ## Objective
#         You are an intelligent and experienced restructuring agent.
#         Your task is to review the already created user journey content {content}.
#         Then, based on the user feedback — "{feedback}", update and restructure the user journey accordingly.
        
#         ## Guidelines
#         - Read and understand the original user journey content.
#         - Understand the user feedback carefully.
#         - If the feedback asks for less content, reduce topics proportionally.
#         - If it asks for more depth, expand topics proportionally.
#         - Do not adjust durations.
        
#         ## Additional Instructions
#         - Keep the format and markdown structure identical to the original.
#         - No projects, but you can include case studies.
#         - Return only the updated markdown, no additional explanation.

#         ## Here is the already existing user journey :
#         {content}
        
#     The updated user journey markdown content as per feedback.

#     """

#     company_context: str = dspy.InputField()
#     content: str = dspy.InputField()
#     feedback: str = dspy.InputField()
#     updated_content: str = dspy.OutputField(
#         desc="The updated user journey markdown syllabus"
#     )


# # === DSPy Module ===
# class RestructureModule(dspy.Module):
#     def __init__(self):
#         super().__init__()
#         self.predictor = dspy.Predict(RestructureSignature)

#         # --- Few-shot examples ---
#         trainset = [
#             dspy.Example(
#                 company_context="""
#                     Company Name: FinSight, 
#                     Industry: Finance, 
#                     Company Size: 250, 
#                     Business Model: B2B, 
#                     Top Use Cases: Risk modeling, 
#                     Learning Focus: Data Analysis, 
#                     Tech Stack: SQL, Tableau, 
#                     Compliance: SOX
#                 """,
#                 content=""" 
#                 ## User Journey\n### Week 1: SQL Basics (3h)\n### Week 2: Data Visualization (3h)""",
#                 feedback="Add more depth on SQL before moving to visualization",
#                 updated_content="""## User Journey\n### Week 1: SQL Basics (3h)\n### Week 2: Advanced SQL Queries (3h)\n### Week 3: Data Visualization (3h)"""
#             ).with_inputs("company_context", "content", "feedback"),
#             dspy.Example(
#                 company_context="""
#                     Company Name: HealthAI, 
#                     Industry: Healthcare, 
#                     Company Size: 500, 
#                     Business Model: SaaS, 
#                     Top Use Cases: Patient analytics, 
#                     Learning Focus: Machine Learning, 
#                     Tech Stack: Python, AWS, 
#                     Compliance: HIPAA
#                 """,
#                 content="""## User Journey\n### Week 1: Python Basics (4h)\n### Week 2: Intro to ML (4h)\n### Week 3: Case Study on Patient Analytics (3h)""",
#                 feedback="First iteration",
#                 updated_content="""## User Journey\n### Week 1: Python Basics (4h)\n### Week 2: Data Handling in Healthcare (4h)\n### Week 3: Intro to ML (4h)\n### Week 4: Case Study on Patient Analytics (3h)"""
#             ).with_inputs("company_context", "content", "feedback"),
#         ]

#         # --- Metric: ensure markdown structure with '## User Journey' ---
#         def restructure_metric(gold, pred, trace=None):
#             output = getattr(pred, "updated_content", "") or ""
#             return int(output.strip().startswith("## User Journey"))

#         optimizer = dspy.BootstrapFewShot(metric=restructure_metric)
#         self.optimized_predictor = optimizer.compile(
#             student=self.predictor,
#             trainset=trainset
#         )

#     def forward(self, company_context, content, feedback):
#         return self.optimized_predictor(
#             company_context=company_context,
#             content=content,
#             feedback=feedback
#         )


# # === Tool Input Schema ===
# class RestructureInput(BaseModel):
#     """Input schema for RestructureTool."""
#     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")
#     content: str = Field(..., description="Existing user journey content in markdown")
#     feedback: str = Field(..., description="User feedback for updating the journey")


# # === Custom DSPy Tool ===
# class RestructureTool(BaseTool):
#     name: str = "Restructure Tool"
#     description: str = (
#         "Restructures or updates a user journey syllabus based on content and user feedback."
#     )
#     args_schema: Type[BaseModel] = RestructureInput

#     def _run(
#         self,
#         Company_Name: str,
#         Industry: str,
#         Company_Size: str,
#         Business_Model: str,
#         Top_Use_Cases: str,
#         Learning_Focus: str,
#         Tech_stack: str,
#         Compliance: str,
#         content: str,
#         feedback: str
#     ) -> str:

#         # Build company context string
#         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}
#         """

#         restructure_module = RestructureModule()
#         result = restructure_module.forward(
#             company_context=company_context,
#             content=content,
#             feedback=feedback
#         )

#         return result.updated_content


# # # === CLI Test Usage ===
# # if __name__ == "__main__":
# #     tool = RestructureTool()
# #     output = tool._run(
# #         Company_Name="DeepCore AI",
# #         Industry="Technology",
# #         Company_Size="500 employees",
# #         Business_Model="SaaS",
# #         Top_Use_Cases="Predictive analytics",
# #         Learning_Focus="Machine Learning",
# #         Tech_stack="Python, AWS, PostgreSQL",
# #         Compliance="SOC 2",
# #         content="## User Journey\n### Week 1: Python Basics (3h)\n### Week 2: Regression Analysis (3h)",
# #         feedback="Add more depth on regression"
# #     )
# #     print("=== Updated User Journey ===")
# #     print(output)
