"""
Baymax Web UI - Minimal Flask interface for case synthesis tools.
Run with: python web_app.py
Access at: http://127.0.0.1:5000
"""

import os
import json
import tempfile
from dataclasses import asdict, is_dataclass
from enum import Enum
from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename


def make_json_serializable(obj):
    """Convert dataclasses and enums to JSON-serializable dicts."""
    if is_dataclass(obj) and not isinstance(obj, type):
        return {k: make_json_serializable(v) for k, v in asdict(obj).items()}
    elif isinstance(obj, Enum):
        return obj.name
    elif isinstance(obj, list):
        return [make_json_serializable(item) for item in obj]
    elif isinstance(obj, dict):
        return {k: make_json_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, tuple):
        return [make_json_serializable(item) for item in obj]
    return obj


# Import Baymax extraction tools
from brief_synthesis import (
    read_case_file,
    synthesize_rules_across_cases,
    generate_parentheticals_for_cases,
    build_citation_string
)
from enhanced_rules import extract_enhanced_rules

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 50MB max upload
app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()

ALLOWED_EXTENSIONS = {'txt', 'pdf', 'docx'}


def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


@app.route('/')
def index():
    """Main page with upload interface."""
    return render_template('index.html')


@app.route('/api/analyze', methods=['POST'])
def analyze_cases():
    """
    Analyze uploaded case files and return synthesized rules.
    Accepts multipart form data with case files.
    """
    if 'files[]' not in request.files:
        return jsonify({'error': 'No files uploaded'}), 400

    files = request.files.getlist('files[]')
    if not files or all(f.filename == '' for f in files):
        return jsonify({'error': 'No files selected'}), 400

    # Save uploaded files temporarily
    saved_paths = []
    case_data = {}  # case_name -> (text, citation) as expected by synthesize functions
    case_texts_only = {}  # case_name -> text for functions that only need text

    try:
        for file in files:
            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
                file.save(filepath)
                saved_paths.append(filepath)

                # Read case content - returns tuple (text, file_type)
                text, file_type = read_case_file(filepath)
                if text and not text.startswith("Error:"):
                    # Extract case name from filename
                    case_name = os.path.splitext(filename)[0]
                    case_name = case_name.replace('_', ' ').replace('-', ' ')
                    # Store as (text, citation) tuple - citation placeholder for now
                    case_data[case_name] = (text, "")
                    case_texts_only[case_name] = text

        if not case_data:
            return jsonify({'error': 'No valid case files could be read'}), 400

        # Extract rules from each case using analyze_rules_deep logic
        all_rules = {}
        for case_name, text in case_texts_only.items():
            rules_result = extract_enhanced_rules(text, case_name)
            if rules_result:
                all_rules[case_name] = rules_result

        # Synthesize rules across cases - expects Dict[str, Tuple[str, str]]
        synthesis_result = synthesize_rules_across_cases(case_data)

        # Generate parentheticals - check what format it expects
        parentheticals = generate_parentheticals_for_cases(case_data)

        # Build response - convert dataclasses/enums to JSON-serializable
        response = make_json_serializable({
            'success': True,
            'cases_analyzed': len(case_data),
            'case_names': list(case_data.keys()),
            'synthesis': synthesis_result,
            'parentheticals': parentheticals,
            'individual_rules': all_rules
        })

        return jsonify(response)

    except Exception as e:
        return jsonify({'error': str(e)}), 500

    finally:
        # Clean up temp files
        for path in saved_paths:
            try:
                os.remove(path)
            except:
                pass


@app.route('/api/build-citation', methods=['POST'])
def build_citation():
    """Build a properly formatted citation string."""
    data = request.get_json()
    if not data:
        return jsonify({'error': 'No data provided'}), 400

    case_name = data.get('case_name', '')
    citation = data.get('citation', '')
    pin_cite = data.get('pin_cite', '')
    parenthetical = data.get('parenthetical', '')
    signal = data.get('signal', '')

    result = build_citation_string(
        case_name=case_name,
        citation=citation,
        pin_cite=pin_cite,
        parenthetical=parenthetical,
        signal=signal
    )

    return jsonify({'citation': result})


@app.route('/api/export', methods=['POST'])
def export_docx():
    """Export synthesis results to Word document."""
    try:
        from docx import Document
        from docx.shared import Pt, Inches
        from docx.enum.text import WD_ALIGN_PARAGRAPH
    except ImportError:
        return jsonify({'error': 'python-docx not installed. Run: pip install python-docx'}), 500

    data = request.get_json()
    if not data:
        return jsonify({'error': 'No data provided'}), 400

    # Create Word document
    doc = Document()

    # Title
    title = doc.add_heading('Rule Synthesis', 0)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER

    # Cases analyzed
    doc.add_heading('Cases Analyzed', level=1)
    for case_name in data.get('case_names', []):
        doc.add_paragraph(case_name, style='List Bullet')

    # Synthesized rules
    if data.get('synthesis'):
        doc.add_heading('Synthesized Rule Hierarchy', level=1)
        synthesis = data['synthesis']

        if isinstance(synthesis, dict):
            if synthesis.get('master_rule'):
                doc.add_heading('Master Rule', level=2)
                doc.add_paragraph(synthesis['master_rule'])

            if synthesis.get('elements'):
                doc.add_heading('Elements', level=2)
                for elem in synthesis['elements']:
                    doc.add_paragraph(elem, style='List Number')

    # Parentheticals
    if data.get('parentheticals'):
        doc.add_heading('Case Parentheticals', level=1)
        for case_name, paren in data['parentheticals'].items():
            p = doc.add_paragraph()
            p.add_run(f"{case_name}: ").bold = True
            p.add_run(paren)

    # Save to temp file
    temp_path = os.path.join(tempfile.gettempdir(), 'baymax_synthesis.docx')
    doc.save(temp_path)

    return send_file(
        temp_path,
        as_attachment=True,
        download_name='rule_synthesis.docx',
        mimetype='application/vnd.openxmlformats-officedocument.wordprocessingml.document'
    )


if __name__ == '__main__':
    # Create templates folder if needed
    templates_dir = os.path.join(os.path.dirname(__file__), 'templates')
    os.makedirs(templates_dir, exist_ok=True)

    print("\n" + "="*50)
    print("  BAYMAX Web UI")
    print("  Open browser to: http://127.0.0.1:5000")
    print("="*50 + "\n")

    app.run(debug=True, port=5000)
