# ai_legal_search_routes.py
from flask import Blueprint, render_template, request, jsonify, session
import requests
import re
import os
import json
from datetime import datetime
from auth_routes import login_required, log_activity

# Create Blueprint
ai_legal_search_bp = Blueprint('ai_legal_search', __name__)

class LegalSearchProcessor:
    def __init__(self, anthropic_api_key=None, openai_api_key=None):
        """Initialize with Anthropic API (preferred) or OpenAI as fallback"""
        self.anthropic_api_key = anthropic_api_key
        self.openai_api_key = openai_api_key
        
        # Prefer Anthropic if available
        if anthropic_api_key:
            self.ai_provider = "anthropic"
            print("✅ Using Anthropic Claude for AI legal research")
        elif openai_api_key:
            import openai
            self.openai_client = openai.OpenAI(api_key=openai_api_key)
            self.ai_provider = "openai"
            print("✅ Using OpenAI GPT-4 for AI legal research")
        else:
            self.ai_provider = None
            print("⚠️ No AI API key found - AI features disabled")
    
    def call_anthropic_api(self, messages, max_tokens=1000):
        """Call Anthropic Claude API"""
        try:
            print(f"🔵 Calling Anthropic API with max_tokens={max_tokens}")

            headers = {
                'Content-Type': 'application/json',
                'x-api-key': self.anthropic_api_key,
                'anthropic-version': '2023-06-01'
            }

            # Convert messages to Anthropic format
            if len(messages) > 1 and messages[0]['role'] == 'system':
                system_message = messages[0]['content']
                user_message = messages[1]['content']
            else:
                system_message = "You are a helpful legal research assistant."
                user_message = messages[0]['content']

            print(f"🔵 System message length: {len(system_message)}, User message length: {len(user_message)}")

            data = {
                'model': 'claude-3-5-sonnet-20241022',
                'max_tokens': max_tokens,
                'system': system_message,
                'messages': [{'role': 'user', 'content': user_message}]
            }

            print("🔵 Sending request to Anthropic...")
            response = requests.post(
                'https://api.anthropic.com/v1/messages',
                headers=headers,
                json=data,
                timeout=30
            )

            print(f"🔵 Anthropic response status: {response.status_code}")

            if response.status_code == 200:
                result = response.json()
                response_text = result['content'][0]['text']
                print(f"🔵 Anthropic response length: {len(response_text)}")
                return response_text
            else:
                print(f"❌ Anthropic API error: {response.status_code} - {response.text}")
                return None

        except Exception as e:
            print(f"❌ Error calling Anthropic API: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def call_ai_api(self, messages, max_tokens=1000):
        """Call AI API (Anthropic preferred, OpenAI fallback)"""
        if self.ai_provider == "anthropic":
            return self.call_anthropic_api(messages, max_tokens)
        elif self.ai_provider == "openai":
            try:
                response = self.openai_client.chat.completions.create(
                    model="gpt-4",
                    messages=messages,
                    max_tokens=max_tokens,
                    temperature=0.3
                )
                return response.choices[0].message.content
            except Exception as e:
                print(f"Error calling OpenAI API: {e}")
                return None
        else:
            return None
        
    def analyze_query(self, user_query):
        """Use AI to analyze the user's query and extract legal search parameters"""
        
        system_prompt = """You are a legal research assistant AI. Analyze the user's natural language query and extract key legal research parameters.

Your task is to identify:
1. Legal topic/area (e.g., "divorce law", "contract law", "employment law")
2. Jurisdiction (state, federal, or specific court if mentioned)
3. Specific legal concepts or statutes mentioned
4. Type of legal authority sought (statutes, case law, regulations)
5. Key search terms for legal databases

Return your analysis as JSON with these fields:
- "topic": Main legal area
- "jurisdiction": State/federal jurisdiction (if specified)
- "legal_concepts": Array of specific legal concepts
- "authority_type": "statutes", "cases", "regulations", or "all"
- "search_terms": Array of optimized search terms
- "summary": Brief summary of what the user is looking for

Example user query: "Nevada statutes for contested divorce proceedings"
Example response:
{
    "topic": "divorce law",
    "jurisdiction": "Nevada",
    "legal_concepts": ["contested divorce", "divorce proceedings"],
    "authority_type": "statutes",
    "search_terms": ["Nevada", "divorce", "contested", "statutes", "proceedings"],
    "summary": "User is looking for Nevada state statutes related to contested divorce proceedings"
}

Respond with ONLY the JSON object, no other text."""

        try:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_query}
            ]
            
            response_text = self.call_ai_api(messages, max_tokens=500)
            
            if response_text:
                # Clean up response to extract JSON
                response_text = response_text.strip()
                if response_text.startswith('```json'):
                    response_text = response_text.replace('```json', '').replace('```', '').strip()
                
                analysis = json.loads(response_text)
                return analysis
            else:
                raise Exception("AI API call failed")
            
        except Exception as e:
            print(f"Error analyzing query: {e}")
            # Fallback analysis
            return {
                "topic": "general legal research",
                "jurisdiction": "general",
                "legal_concepts": [user_query],
                "authority_type": "all",
                "search_terms": user_query.split(),
                "summary": f"User is looking for information about: {user_query}"
            }

    def search_legal_databases(self, analysis):
        """Search various legal databases based on the analysis"""
        results = []
        
        # Search CourtListener for case law
        if analysis["authority_type"] in ["cases", "all"]:
            case_results = self.search_courtlistener(analysis)
            results.extend(case_results)
        
        # Search for statutes
        if analysis["authority_type"] in ["statutes", "all"]:
            statute_results = self.search_statutes(analysis)
            results.extend(statute_results)
        
        # Search for regulations
        if analysis["authority_type"] in ["regulations", "all"]:
            regulation_results = self.search_regulations(analysis)
            results.extend(regulation_results)
            
        return results

    def search_courtlistener(self, analysis):
        """Search CourtListener for relevant case law"""
        results = []
        
        try:
            api_token = os.environ.get('COURTLISTENER_API_TOKEN')
            if not api_token:
                return results
                
            headers = {
                'Authorization': f'Token {api_token}',
                'User-Agent': 'LawBot AI Legal Research'
            }
            
            # Build search query
            search_terms = " ".join(analysis["search_terms"][:5])  # Limit terms
            
            search_params = {
                'q': search_terms,
                'type': 'o',  # Opinions
                'order_by': 'score desc'
            }
            
            # Add jurisdiction filter if specified
            if analysis["jurisdiction"] and analysis["jurisdiction"].lower() != "general":
                search_params['court'] = analysis["jurisdiction"].lower()
            
            response = requests.get(
                'https://www.courtlistener.com/api/rest/v4/search/',
                params=search_params,
                headers=headers,
                timeout=10
            )
            
            if response.status_code == 200:
                data = response.json()
                
                for item in data.get('results', [])[:5]:  # Limit to top 5 results
                    result = {
                        'type': 'case',
                        'title': self.extract_case_name(item),
                        'citation': self.extract_citation(item),
                        'description': self.extract_case_synopsis(item),
                        'url': item.get('absolute_url', ''),
                        'court': self.extract_court_name(item),
                        'date': item.get('dateFiled', item.get('date_filed', ''))
                    }
                    
                    if result['url'] and not result['url'].startswith('http'):
                        result['url'] = f"https://www.courtlistener.com{result['url']}"
                    
                    results.append(result)
                    
        except Exception as e:
            print(f"Error searching CourtListener: {e}")
            
        return results

    def search_statutes(self, analysis):
        """Search for relevant statutes"""
        results = []
        
        jurisdiction = analysis.get("jurisdiction", "").lower()
        topic = analysis.get("topic", "").lower()
        
        statute_mappings = {
            "divorce": {
                "nevada": [
                    {
                        'type': 'statute',
                        'title': 'Nevada Revised Statutes Chapter 125 - Dissolution of Marriage',
                        'citation': 'NRS 125.010 et seq.',
                        'description': 'Nevada statutes governing divorce proceedings, including grounds for divorce, property division, and child custody.',
                        'url': 'https://www.leg.state.nv.us/nrs/nrs-125.html'
                    },
                    {
                        'type': 'statute',
                        'title': 'Nevada Revised Statutes 125.150 - Contested Proceedings',
                        'citation': 'NRS 125.150',
                        'description': 'Specific statute addressing contested divorce proceedings in Nevada.',
                        'url': 'https://www.leg.state.nv.us/nrs/nrs-125.html#NRS125Sec150'
                    }
                ],
                "california": [
                    {
                        'type': 'statute',
                        'title': 'California Family Code Division 6 - Nullity, Dissolution, and Legal Separation',
                        'citation': 'Cal. Fam. Code § 2100 et seq.',
                        'description': 'California statutes governing divorce and dissolution of marriage.',
                        'url': 'https://leginfo.legislature.ca.gov/faces/codes_displayexpandedbranch.xhtml?tocCode=FAM&division=6'
                    }
                ]
            },
            "landlord tenant": {
                "california": [
                    {
                        'type': 'statute',
                        'title': 'California Civil Code Section 1940 et seq. - Landlord-Tenant Law',
                        'citation': 'Cal. Civ. Code § 1940 et seq.',
                        'description': 'California statutes governing landlord-tenant relationships, including eviction procedures.',
                        'url': 'https://leginfo.legislature.ca.gov/faces/codes_displaySection.xhtml?sectionNum=1940&lawCode=CIV'
                    }
                ]
            },
            "employment": {
                "federal": [
                    {
                        'type': 'statute',
                        'title': 'Title VII of the Civil Rights Act of 1964',
                        'citation': '42 U.S.C. § 2000e et seq.',
                        'description': 'Federal law prohibiting employment discrimination based on race, color, religion, sex, or national origin.',
                        'url': 'https://www.eeoc.gov/statutes/title-vii-civil-rights-act-1964'
                    },
                    {
                        'type': 'statute',
                        'title': 'Fair Labor Standards Act',
                        'citation': '29 U.S.C. § 201 et seq.',
                        'description': 'Federal law establishing minimum wage, overtime pay, and child labor standards.',
                        'url': 'https://www.dol.gov/agencies/whd/flsa'
                    }
                ]
            },
            "contract": {
                "general": [
                    {
                        'type': 'statute',
                        'title': 'Uniform Commercial Code Article 2 - Sales',
                        'citation': 'UCC § 2-101 et seq.',
                        'description': 'Uniform laws governing sales of goods and contract remedies.',
                        'url': 'https://www.law.cornell.edu/ucc/2'
                    }
                ]
            },
            "criminal": {
                "federal": [
                    {
                        'type': 'statute',
                        'title': 'Federal Criminal Code',
                        'citation': '18 U.S.C. § 1 et seq.',
                        'description': 'Federal criminal statutes and procedures.',
                        'url': 'https://www.law.cornell.edu/uscode/text/18'
                    }
                ]
            },
            "immigration": {
                "federal": [
                    {
                        'type': 'statute',
                        'title': 'Immigration and Nationality Act',
                        'citation': '8 U.S.C. § 1101 et seq.',
                        'description': 'Federal immigration laws and visa requirements.',
                        'url': 'https://www.uscis.gov/laws-and-policy/immigration-and-nationality-act'
                    }
                ]
            },
            "bankruptcy": {
                "federal": [
                    {
                        'type': 'statute',
                        'title': 'Bankruptcy Code Title 11',
                        'citation': '11 U.S.C. § 101 et seq.',
                        'description': 'Federal bankruptcy laws governing liquidation and reorganization proceedings.',
                        'url': 'https://www.law.cornell.edu/uscode/text/11'
                    }
                ]
            }
        }
        
        # Try to match topic and jurisdiction
        for topic_key in statute_mappings:
            if topic_key in topic:
                if jurisdiction in statute_mappings[topic_key]:
                    results.extend(statute_mappings[topic_key][jurisdiction])
                elif "federal" in statute_mappings[topic_key]:
                    results.extend(statute_mappings[topic_key]["federal"])
                elif "general" in statute_mappings[topic_key]:
                    results.extend(statute_mappings[topic_key]["general"])
                    
        return results

    def search_regulations(self, analysis):
        """Search for relevant regulations"""
        results = []
        
        topic = analysis.get("topic", "").lower()
        
        if "employment" in topic:
            results.append({
                'type': 'regulation',
                'title': 'Equal Employment Opportunity Regulations',
                'citation': '29 CFR Part 1600 et seq.',
                'description': 'Federal regulations implementing employment discrimination laws.',
                'url': 'https://www.ecfr.gov/current/title-29/subtitle-B/chapter-XIV'
            })
        
        if "immigration" in topic:
            results.append({
                'type': 'regulation',
                'title': 'Immigration Regulations',
                'citation': '8 CFR et seq.',
                'description': 'Federal regulations implementing immigration and nationality laws.',
                'url': 'https://www.ecfr.gov/current/title-8'
            })
            
        return results

    def extract_case_name(self, case_data):
        """Extract case name from API response"""
        case_name_fields = ['caseName', 'caseNameFull', 'case_name', 'title']
        for field in case_name_fields:
            if field in case_data and case_data[field]:
                return str(case_data[field]).strip()
        return 'Untitled Case'

    def extract_citation(self, case_data):
        """Extract citation from API response"""
        if 'neutralCite' in case_data and case_data['neutralCite']:
            return case_data['neutralCite']
        if 'citation' in case_data and case_data['citation']:
            return str(case_data['citation'])
        return 'Citation not available'

    def extract_case_synopsis(self, case_data):
        """Extract case synopsis from API response"""
        synopsis_fields = ['syllabus', 'snippet', 'summary']
        for field in synopsis_fields:
            if field in case_data and case_data[field]:
                synopsis = str(case_data[field]).strip()
                if len(synopsis) > 20:
                    return synopsis[:200] + '...' if len(synopsis) > 200 else synopsis
        return 'Synopsis not available'

    def extract_court_name(self, case_data):
        """Extract court name from API response"""
        if 'court' in case_data and case_data['court']:
            if isinstance(case_data['court'], dict):
                return case_data['court'].get('name', case_data['court'].get('short_name', 'Unknown Court'))
            return str(case_data['court'])
        return 'Unknown Court'

    def generate_response(self, user_query, analysis, search_results):
        """Generate a natural language response using AI"""

        print(f"🟢 generate_response() called with {len(search_results)} search results")

        # Prepare context for AI response
        context = f"""
User Query: {user_query}
Legal Topic: {analysis['topic']}
Jurisdiction: {analysis['jurisdiction']}
Search Summary: {analysis['summary']}

Found {len(search_results)} relevant legal sources.
"""
        print(f"🟢 Context prepared, length: {len(context)}")
        
        system_prompt = """You are a helpful legal research assistant. Based on the user's query and the search results found, provide a clear, informative response that:

1. Directly addresses the user's question
2. Summarizes the relevant legal authorities found
3. Explains how the sources relate to their query
4. Provides practical guidance on next steps
5. Maintains a professional but accessible tone

Do not provide specific legal advice. Always remind users to consult with qualified legal counsel for specific situations.

Keep your response concise but informative (2-3 paragraphs maximum).

Important: Do not mention or reference specific case names, citations, or detailed legal content from the search results in your response. Instead, provide general guidance and direct the user to review the sources provided."""

        try:
            # Create a summary of search results for the AI
            results_summary = []
            for result in search_results[:10]:  # Limit to top 10
                results_summary.append(f"- {result['title']} ({result.get('citation', 'No citation')}): {result.get('description', 'No description')}")
            
            results_text = "\n".join(results_summary) if results_summary else "No specific sources found."
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"{context}\n\nRelevant Sources Found:\n{results_text}\n\nPlease provide a helpful response to the user's query."}
            ]
            
            print(f"🟢 Calling AI API for response generation...")
            response_text = self.call_ai_api(messages, max_tokens=600)

            print(f"🟢 AI API returned: {type(response_text)}, Length: {len(response_text) if response_text else 0}")

            if response_text:
                print(f"🟢 Returning response text")
                return response_text
            else:
                print(f"❌ AI API call returned None/empty")
                raise Exception("AI API call failed")

        except Exception as e:
            print(f"❌ Error generating response: {e}")
            import traceback
            traceback.print_exc()

            # Fallback response
            if search_results:
                fallback = f"I found {len(search_results)} relevant legal sources for your query about {analysis['topic']}. Please review the sources below for detailed information. For specific legal advice regarding your situation, please consult with a qualified attorney."
                print(f"🟡 Returning fallback response (with results): {fallback[:100]}")
                return fallback
            else:
                fallback = f"I wasn't able to find specific legal sources for your query about {analysis['topic']}. You may want to try refining your search terms or consulting with a legal professional for assistance with this topic."
                print(f"🟡 Returning fallback response (no results): {fallback[:100]}")
                return fallback

# Initialize the processor
legal_search_processor = None

def init_legal_search(app):
    """Initialize the legal search processor with API keys from app config"""
    global legal_search_processor
    
    anthropic_api_key = app.config.get('ANTHROPIC_API_KEY') or os.environ.get('ANTHROPIC_API_KEY')
    openai_api_key = app.config.get('OPENAI_API_KEY') or os.environ.get('OPENAI_API_KEY')
    
    if anthropic_api_key or openai_api_key:
        legal_search_processor = LegalSearchProcessor(
            anthropic_api_key=anthropic_api_key,
            openai_api_key=openai_api_key
        )
        print("✅ AI Legal Search initialized successfully")
    else:
        print("⚠️ Warning: No AI API keys found, AI Legal Search disabled")

# Routes
@ai_legal_search_bp.route('/ai-research')
@login_required
def ai_research_page():
    """Display the AI legal research interface"""
    return render_template('ai_legal_research.html')

@ai_legal_search_bp.route('/api/ai-legal-search', methods=['POST'])
@login_required
def ai_legal_search():
    """Process natural language legal research queries"""
    try:
        if not legal_search_processor:
            return jsonify({
                'success': False,
                'error': 'AI Legal Search is not available'
            }), 500

        data = request.get_json()
        user_query = data.get('query', '').strip()
        
        if not user_query:
            return jsonify({
                'success': False,
                'error': 'Query is required'
            }), 400

        print(f"Processing AI legal search query: {user_query}")

        # Step 1: Analyze the query using AI
        analysis = legal_search_processor.analyze_query(user_query)
        print(f"Query analysis: {analysis}")

        # Step 2: Search legal databases
        search_results = legal_search_processor.search_legal_databases(analysis)
        print(f"Found {len(search_results)} search results")

        # Step 3: Generate AI response
        ai_response = legal_search_processor.generate_response(user_query, analysis, search_results)

        # DEBUG: Log response details
        print(f"AI Response generated - Length: {len(ai_response) if ai_response else 0}")
        print(f"AI Response preview: {ai_response[:200] if ai_response else 'EMPTY/NONE'}")

        # Step 4: Log the activity
        user_id = session.get('user_id')
        if user_id:
            log_activity(
                user_id,
                'ai_legal_search',
                f'AI legal search: {user_query[:100]}...' if len(user_query) > 100 else f'AI legal search: {user_query}',
                'search',
                None
            )

        # Ensure we have a response
        if not ai_response or ai_response.strip() == '':
            ai_response = "I apologize, but I wasn't able to generate a response. Please try rephrasing your question or contact support if this issue persists."

        return jsonify({
            'success': True,
            'response': ai_response,
            'sources': search_results,
            'analysis': analysis,
            'query': user_query
        })

    except Exception as e:
        print(f"Error in AI legal search: {e}")
        import traceback
        traceback.print_exc()
        
        return jsonify({
            'success': False,
            'error': 'An error occurred while processing your search. Please try again.'
        }), 500

@ai_legal_search_bp.route('/api/ai-legal-search/suggestions')
@login_required
def get_search_suggestions():
    """Get suggested search queries based on common legal topics"""
    suggestions = [
        "Nevada statutes for contested divorce proceedings",
        "California landlord tenant law eviction process", 
        "Federal employment discrimination laws",
        "Personal injury statute of limitations by state",
        "Contract law breach of contract remedies",
        "Criminal procedure Miranda rights requirements",
        "Bankruptcy Chapter 7 vs Chapter 13 differences",
        "Intellectual property copyright fair use doctrine",
        "Real estate purchase agreement essential terms",
        "Workers compensation claim process",
        "Small claims court filing procedures",
        "Immigration visa requirements and processes"
    ]
    
    return jsonify({
        'success': True,
        'suggestions': suggestions
    })