import streamlit as st
from datetime import datetime
import os
from typing import Dict
from fpdf import FPDF
import anthropic


class TemplateGenerator:
    def __init__(self):
        # Initialize Anthropic client with API key
        try:
            api_key = os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                st.error("Please set the ANTHROPIC_API_KEY in your environment variables.")
                st.stop()
            self.client = anthropic.Anthropic(api_key=api_key)
        except Exception as e:
            st.error(f"Error initializing Anthropic client: {str(e)}")
            st.stop()

        # Define templates
        self.templates = {
            "nda": {
                "name": "Non-Disclosure Agreement",
                "description": "Standard NDA for business confidentiality",
                "fields": {
                    "company_name": "Company requiring NDA",
                    "recipient_name": "Party signing the NDA",
                    "effective_date": "Start date of agreement",
                    "duration_months": "Duration in months",
                    "jurisdiction": "Governing law jurisdiction",
                    "confidential_info_description": "Description of confidential information",
                    "permitted_use": "Allowed uses of confidential information",
                    "return_period": "Days to return/destroy information"
                }
            },
            "services_agreement": {
                "name": "Services Agreement",
                "description": "Comprehensive contract for provision of services",
                "fields": {
                    "service_provider": "Provider company/individual name",
                    "client_name": "Client name",
                    "services_description": "Detailed service description",
                    "start_date": "Service start date",
                    "payment_terms": "Payment amount and schedule",
                    "termination_notice": "Notice period for termination",
                    "governing_law": "Jurisdiction",
                    "service_levels": "Service level agreements (SLAs)"
                }
            },
            # Extend templates here
        }

    def render(self):
        """Render template generator interface."""
        st.header("Legal Document Generator")
        st.write("Generate comprehensive legal documents using customizable templates.")

        # Template selection
        template_type = st.selectbox(
            "Select Document Type",
            options=list(self.templates.keys()),
            format_func=lambda x: self.templates[x]["name"]
        )

        # Show template description
        st.info(self.templates[template_type]["description"])

        # Input form
        with st.form(f"template_form_{template_type}"):
            inputs = {}
            for field, description in self.templates[template_type]["fields"].items():
                if "date" in field:
                    inputs[field] = st.date_input(f"{description}")
                elif "description" in field or "terms" in field:
                    inputs[field] = st.text_area(f"{description}", height=100)
                elif "amount" in field or "months" in field:
                    inputs[field] = st.number_input(f"{description}", min_value=0)
                else:
                    inputs[field] = st.text_input(f"{description}")

            # Submit form
            generate = st.form_submit_button("Generate Document")

        if generate:
            if all(inputs.values()):
                with st.spinner("Generating document..."):
                    document = self._generate_document(template_type, inputs)
                    self._display_document(document, template_type)
            else:
                st.error("Please fill in all required fields.")

    def _generate_document(self, template_type: str, inputs: Dict) -> str:
        """Generate document using Claude."""
        try:
            prompt = self._create_prompt(template_type, inputs)
            message = self.client.messages.create(
                model="claude-3-sonnet-20240229",
                max_tokens=2000,
                temperature=0.7,
                messages=[{"role": "user", "content": prompt}]
            )
            return message.content[0].text
        except Exception as e:
            st.error(f"Error generating document: {str(e)}")
            return ""

    def _create_prompt(self, template_type: str, inputs: Dict) -> str:
        """Create prompt for document generation."""
        template_name = self.templates[template_type]["name"]

        prompt = f"""Generate a detailed {template_name} using the following information:

Template Type: {template_name}
"""
        for field, value in inputs.items():
            prompt += f"{field}: {value}\n"

        prompt += """

Ensure the document:
1. Follows professional legal formatting.
2. Includes all necessary clauses and sections.
3. Is tailored for {inputs.get('jurisdiction', 'general')} jurisdiction.
4. Uses formal legal language.
5. Includes signature blocks and clear numbering.
"""
        return prompt

    def _display_document(self, document: str, template_type: str):
        """Display and provide download options for the generated document."""
        st.subheader("Generated Document")
        with st.expander("Preview", expanded=True):
            st.markdown(document)

        col1, col2 = st.columns(2)
        with col1:
            st.download_button(
                "Download as TXT",
                document,
                file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.txt",
                mime="text/plain"
            )
        with col2:
            pdf_document = self._generate_pdf(document)
            st.download_button(
                "Download as PDF",
                pdf_document,
                file_name=f"{template_type}_{datetime.now().strftime('%Y%m%d')}.pdf",
                mime="application/pdf"
            )

    def _generate_pdf(self, document: str) -> bytes:
        """Generate a PDF file for the document."""
        pdf = FPDF()
        pdf.add_page()
        pdf.set_font("Arial", size=12)
        for line in document.split("\n"):
            pdf.cell(200, 10, txt=line, ln=True)
        output = bytes()
        pdf.output(name="document.pdf", dest="S").encode("latin1")
        return output


def render_template_generator():
    """Main function to render the template generator."""
    generator = TemplateGenerator()
    generator.render()