from cerevox import AsyncLexa
import re
from typing import Dict, List

class ContractAnalyzer:
    def __init__(self, api_key: str):
        self.client = AsyncLexa(api_key=api_key)
        
        # Define clause patterns
        self.clause_patterns = {
            'termination': [
                r'terminat[e|ion].*?(?:\.|;|\n)',
                r'end.*?agreement.*?(?:\.|;|\n)'
            ],
            'liability': [
                r'liabilit.*?(?:\.|;|\n)',
                r'indemnif.*?(?:\.|;|\n)'
            ],
            'confidentiality': [
                r'confidential.*?(?:\.|;|\n)',
                r'proprietary.*?(?:\.|;|\n)'
            ]
        }
    
    async def analyze_contract(self, contract_path: str) -> Dict:
        """Comprehensive contract analysis"""
        async with self.client:
            documents = await self.client.parse(contract_path)
            contract = documents[0]
            
            analysis = {
                'contract_info': {
                    'filename': contract.filename,
                    'pages': contract.total_pages
                },
                'extracted_clauses': {},
                'risk_assessment': self._assess_risks(contract)
            }
            
            # Extract clauses by type
            for clause_type, patterns in self.clause_patterns.items():
                clauses = []
                for pattern in patterns:
                    matches = re.findall(pattern, contract.content, re.IGNORECASE)
                    clauses.extend(matches)
                
                analysis['extracted_clauses'][clause_type] = {
                    'count': len(clauses),
                    'text': clauses[:3]  # First 3 matches
                }
            
            return analysis
    
    def _assess_risks(self, contract) -> Dict:
        """Assess legal risks in contract"""
        high_risk_terms = [
            "unlimited liability", "personal guarantee", 
            "liquidated damages", "penalty"
        ]
        
        risks = {}
        for term in high_risk_terms:
            matches = contract.search_content(term)
            if matches:
                risks[term] = {
                    'count': len(matches),
                    'severity': 'high'
                }
        
        return risks

# Usage
analyzer = ContractAnalyzer("your-api-key")
analysis = await analyzer.analyze_contract("service_agreement.pdf")

from cerevox import AsyncLexa
from typing import List, Dict

class LegalRAGSystem:
    def __init__(self, api_key: str):
        self.client = AsyncLexa(api_key=api_key)
        self.legal_corpus = []
    
    async def ingest_legal_documents(self, document_paths: List[str]):
        """Ingest legal documents for search"""
        async with self.client:
            documents = await self.client.parse(document_paths)
            
            # Create searchable chunks optimized for legal content
            for doc in documents:
                chunks = doc.get_text_chunks(
                    target_size=800,  # Good for legal context
                    tolerance=0.2
                )
                
                for i, chunk in enumerate(chunks):
                    self.legal_corpus.append({
                        'id': f"{doc.filename}_{i}",
                        'content': chunk,
                        'document': doc.filename,
                        'document_type': self._classify_document(doc.filename)
                    })
        
        return f"Ingested {len(self.legal_corpus)} legal document chunks"
    
    def search_legal_precedents(self, query: str, limit: int = 5) -> List[Dict]:
        """Search for legal precedents and relevant cases"""
        results = []
        query_terms = query.lower().split()
        
        for chunk in self.legal_corpus:
            content_lower = chunk['content'].lower()
            
            # Score based on term frequency
            score = sum(1 for term in query_terms if term in content_lower)
            
            if score > 0:
                results.append({
                    'content': chunk['content'],
                    'document': chunk['document'],
                    'relevance_score': score,
                    'document_type': chunk['document_type']
                })
        
        # Sort by relevance and return top results
        results.sort(key=lambda x: x['relevance_score'], reverse=True)
        return results[:limit]
    
    def _classify_document(self, filename: str) -> str:
        """Classify legal document type"""
        filename_lower = filename.lower()
        
        if any(term in filename_lower for term in ['contract', 'agreement']):
            return 'contract'
        elif any(term in filename_lower for term in ['case', 'opinion']):
            return 'case_law'
        elif any(term in filename_lower for term in ['brief', 'motion']):
            return 'court_filing'
        else:
            return 'legal_document'

# Usage
legal_rag = LegalRAGSystem("your-api-key")
await legal_rag.ingest_legal_documents(["contracts.pdf", "case_law.pdf"])
results = legal_rag.search_legal_precedents("intellectual property")

async def automated_due_diligence(target_company_docs):
    """Automate due diligence document review"""
    
    async with AsyncLexa(api_key="your-api-key") as client:
        documents = await client.parse(target_company_docs)
        
        due_diligence_report = {}
        for doc in documents:
            # Extract critical information
            due_diligence_report[doc.filename] = {
                'contract_terms': doc.search_content("term|clause|provision"),
                'liabilities': doc.search_content("liabilit|debt|obligation"),
                'intellectual_property': doc.search_content("patent|trademark"),
                'litigation_risks': doc.search_content("lawsuit|litigation"),
                'compliance_issues': doc.search_content("regulation|compliance")
            }
        
        return due_diligence_report

async def process_discovery_documents(document_list: List[str]):
    """Process large volumes of discovery documents"""
    
    async with AsyncLexa(api_key="your-api-key") as client:
        BATCH_SIZE = 10
        all_results = {}
        
        for i in range(0, len(document_list), BATCH_SIZE):
            batch = document_list[i:i + BATCH_SIZE]
            documents = await client.parse(batch)
            
            for doc in documents:
                key_info = {
                    'pages': doc.total_pages,
                    'communications': doc.search_content("email|letter|memo"),
                    'privileged_content': doc.search_content("attorney|counsel"),
                    'key_people': extract_people_names(doc)
                }
                
                all_results[doc.filename] = key_info
        
        return all_results