#!/usr/bin/env python3
"""
IT Site Survey AI Flask Application
Creates surveys for client sites and uses AI to recommend network setups.
Enhanced with PDF export and quote generation.
"""

import os
import json
import uuid
import requests
from datetime import datetime
from flask import Flask, jsonify, request, send_from_directory, Response, stream_with_context, send_file
from flask_cors import CORS
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, ListFlowable, ListItem
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
import io

app = Flask(__name__, static_url_path='', static_folder='.')
CORS(app)

OLLAMA_BASE = os.environ.get('OLLAMA_URL', 'http://192.168.19.25:11434')

# In-memory storage for surveys
surveys_db = {}
survey_responses_db = {}

# Default survey template
DEFAULT_SURVEY_TEMPLATE = {
  "id": "concise-network-survey-template",
  "name": "Concise Network Site Survey",
  "description": "Streamlined one-page network infrastructure survey with dropdown selections",
  "questions": [
    {
      "id": "company_name",
      "category": "Client Information",
      "question": "Company Name",
      "type": "text"
    },
    {
      "id": "industry",
      "category": "Client Information",
      "question": "Industry/Vertical",
      "type": "select",
      "options": [
        "Healthcare",
        "Education",
        "Financial Services",
        "Manufacturing",
        "Retail",
        "Technology",
        "Government",
        "Hospitality",
        "Other"
      ]
    },
    {
      "id": "site_size",
      "category": "Site Information",
      "question": "Site Size",
      "type": "select",
      "options": [
        "Small (1-10 employees)",
        "Medium (11-50 employees)",
        "Large (51-250 employees)",
        "Enterprise (250+ employees)"
      ]
    },
    {
      "id": "building_type",
      "category": "Site Information",
      "question": "Building Type",
      "type": "select",
      "options": [
        "Single Office",
        "Multi-tenant",
        "Campus",
        "Warehouse",
        "Retail",
        "Industrial",
        "Data Center",
        "Other"
      ]
    },
    {
      "id": "square_footage",
      "category": "Site Information",
      "question": "Total Square Footage",
      "type": "select",
      "options": [
        "Under 5,000 sq ft",
        "5,000 - 10,000 sq ft",
        "10,000 - 25,000 sq ft",
        "25,000 - 50,000 sq ft",
        "Over 50,000 sq ft"
      ]
    },
    {
      "id": "network_topology",
      "category": "Existing Infrastructure",
      "question": "Current Network Topology",
      "type": "select",
      "options": [
        "None - New Installation",
        "Flat Network",
        "Two-tier (Access/Distribution)",
        "Three-tier (Access/Distribution/Core)",
        "Spine-Leaf",
        "Other/Not Sure"
      ]
    },
    {
      "id": "network_size",
      "category": "Existing Infrastructure",
      "question": "Current Network Size",
      "type": "select",
      "options": [
        "Small (<50 devices)",
        "Medium (50-250 devices)",
        "Large (250-1000 devices)",
        "Enterprise (1000+ devices)"
      ]
    },
    {
      "id": "internet_bandwidth",
      "category": "Connectivity",
      "question": "Current Internet Bandwidth",
      "type": "select",
      "options": [
        "Under 50 Mbps",
        "50-100 Mbps",
        "100-500 Mbps",
        "500 Mbps-1 Gbps",
        "1-10 Gbps",
        "Over 10 Gbps",
        "Not Sure"
      ]
    },
    {
      "id": "isp_type",
      "category": "Connectivity",
      "question": "ISP Connection Type",
      "type": "multiselect",
      "options": [
        "Fiber",
        "Cable/DSL",
        "Fixed Wireless",
        "5G/LTE",
        "Multiple ISPs",
        "Satellite",
        "Not Sure"
      ]
    },
    {
      "id": "wireless_standard",
      "category": "Wireless Infrastructure",
      "question": "Current Wireless Standard",
      "type": "select",
      "options": [
        "None",
        "802.11n (Wi-Fi 4)",
        "802.11ac (Wi-Fi 5)",
        "802.11ax (Wi-Fi 6)",
        "802.11be (Wi-Fi 7)",
        "Mixed/Not Sure"
      ]
    },
    {
      "id": "wireless_coverage",
      "category": "Wireless Infrastructure",
      "question": "Wireless Coverage Quality",
      "type": "select",
      "options": [
        "Excellent - No dead zones",
        "Good - Minor coverage issues",
        "Fair - Several dead zones",
        "Poor - Major coverage issues",
        "No wireless network"
      ]
    },
    {
      "id": "critical_applications",
      "category": "Business Requirements",
      "question": "Critical Business Applications",
      "type": "multiselect",
      "options": [
        "Email/Collaboration (Office 365, Google Workspace)",
        "VoIP/Video Conferencing",
        "Cloud Applications",
        "File Sharing/Storage",
        "Database Applications",
        "Remote Desktop/VDI",
        "Video Surveillance",
        "Point of Sale (POS)",
        "Industrial Systems",
        "Other"
      ]
    },
    {
      "id": "device_types",
      "category": "Device Requirements",
      "question": "Primary Device Types",
      "type": "multiselect",
      "options": [
        "Desktop Computers",
        "Laptops",
        "Mobile Devices (Phones/Tablets)",
        "IoT Devices",
        "Printers/MFPs",
        "Servers",
        "Video Equipment",
        "Specialized Equipment"
      ]
    },
    {
      "id": "user_density",
      "category": "Usage Requirements",
      "question": "User Density",
      "type": "select",
      "options": [
        "Low (<10 devices per AP)",
        "Medium (10-30 devices per AP)",
        "High (30-50 devices per AP)",
        "Very High (50+ devices per AP)"
      ]
    },
    {
      "id": "performance_requirements",
      "category": "Performance Requirements",
      "question": "Network Performance Requirements",
      "type": "multiselect",
      "options": [
        "High Speed (1+ Gbps)",
        "Low Latency (<10ms)",
        "High Availability (99.9%+)",
        "Guest Network Access",
        "BYOD Support",
        "Outdoor Coverage",
        "No Special Requirements"
      ]
    },
    {
      "id": "security_requirements",
      "category": "Security Requirements",
      "question": "Security Requirements",
      "type": "multiselect",
      "options": [
        "Basic Firewall",
        "Content Filtering",
        "Network Segmentation",
        "802.1X Authentication",
        "Guest Network Isolation",
        "Compliance (HIPAA, PCI, etc.)",
        "No Special Security Requirements"
      ]
    },
    {
      "id": "compliance_needs",
      "category": "Compliance Requirements",
      "question": "Compliance Requirements",
      "type": "select",
      "options": [
        "None",
        "HIPAA (Healthcare)",
        "PCI-DSS (Payment Card)",
        "SOX (Financial)",
        "FERPA (Educational)",
        "GDPR (Privacy)",
        "Industry Specific",
        "Multiple Compliance Requirements"
      ]
    },
    {
      "id": "budget_range",
      "category": "Project Information",
      "question": "Estimated Budget Range",
      "type": "select",
      "options": [
        "Under $5,000",
        "$5,000 - $15,000",
        "$15,000 - $50,000",
        "$50,000 - $100,000",
        "$100,000+",
        "Not Established"
      ]
    },
    {
      "id": "timeline",
      "category": "Project Information",
      "question": "Project Timeline",
      "type": "select",
      "options": [
        "ASAP - Emergency",
        "Within 2 weeks",
        "Within 1 month",
        "Within 3 months",
        "Within 6 months",
        "Flexible/Planning Phase"
      ]
    },
    {
      "id": "additional_notes",
      "category": "Additional Information",
      "question": "Additional Notes or Special Requirements",
      "type": "textarea",
      "placeholder": "Please share any other important details about your network requirements..."
    }
  ]
}

def call_ollama(prompt, model='llama3.2:3b', stream=False):
    """Call Ollama API with the given prompt."""
    try:
        if stream:
            resp = requests.post(
                f'{OLLAMA_BASE}/api/generate',
                json={'model': model, 'prompt': prompt, 'stream': True},
                stream=True,
                timeout=120
            )
            return resp
        else:
            resp = requests.post(
                f'{OLLAMA_BASE}/api/generate',
                json={'model': model, 'prompt': prompt, 'stream': False},
                timeout=120
            )
            return resp.json().get('response', '')
    except Exception as e:
        return f'Error: {str(e)}'


@app.route('/')
def index():
    return send_from_directory('.', 'index.html')


@app.route('/<path:path>')
def serve_static(path):
    return send_from_directory('.', path)


@app.route('/api/health')
def health_check():
    try:
        r = requests.get(f'{OLLAMA_BASE}/', timeout=3)
        ollama_online = r.status_code == 200
    except:
        ollama_online = False
    
    return jsonify({
        'status': 'ok',
        'service': 'site-survey-ai',
        'ollama_connected': ollama_online,
        'timestamp': datetime.utcnow().isoformat()
    })


@app.route('/api/ollama-status')
def ollama_status():
    try:
        r = requests.get(f'{OLLAMA_BASE}/', timeout=3)
        return jsonify({'online': r.status_code == 200})
    except:
        return jsonify({'online': False})


@app.route('/api/surveys/template')
def get_template():
    return jsonify({'template': DEFAULT_SURVEY_TEMPLATE})


@app.route('/api/surveys', methods=['GET'])
def get_surveys():
    return jsonify({
        'surveys': list(surveys_db.values())
    })


@app.route('/api/surveys', methods=['POST'])
def create_survey():
    data = request.get_json()
    
    survey_id = str(uuid.uuid4())
    survey = {
        'id': survey_id,
        'name': data.get('name', 'New Survey'),
        'description': data.get('description', ''),
        'client_name': data.get('client_name', ''),
        'site_name': data.get('site_name', ''),
        'questions': data.get('questions', DEFAULT_SURVEY_TEMPLATE['questions']),
        'created_at': datetime.utcnow().isoformat(),
        'status': 'active',
        'responses': []
    }
    
    surveys_db[survey_id] = survey
    
    return jsonify({'survey': survey}), 201


@app.route('/api/surveys/<survey_id>', methods=['GET'])
def get_survey(survey_id):
    survey = surveys_db.get(survey_id)
    if not survey:
        return jsonify({'error': 'Survey not found'}), 404
    return jsonify({'survey': survey})


@app.route('/api/surveys/<survey_id>/responses', methods=['POST'])
def submit_response(survey_id):
    survey = surveys_db.get(survey_id)
    if not survey:
        return jsonify({'error': 'Survey not found'}), 404
    
    data = request.get_json()
    
    response_id = str(uuid.uuid4())
    response = {
        'id': response_id,
        'survey_id': survey_id,
        'submitted_by': data.get('submitted_by', 'Anonymous'),
        'responses': data.get('responses', {}),
        'submitted_at': datetime.utcnow().isoformat()
    }
    
    survey['responses'].append(response)
    survey_responses_db[response_id] = response
    
    return jsonify({'response': response}), 201


@app.route('/api/surveys/responses', methods=['GET'])
def get_all_responses():
    """Get all survey responses across all surveys."""
    all_responses = list(survey_responses_db.values())
    return jsonify({'responses': all_responses})


@app.route('/api/surveys/<survey_id>/analyze', methods=['POST'])
def analyze_survey(survey_id):
    """Analyze survey responses with AI to generate recommendations."""
    survey = surveys_db.get(survey_id)
    if not survey:
        return jsonify({'error': 'Survey not found'}), 404
    
    if not survey['responses']:
        return jsonify({'error': 'No responses to analyze'}), 400
    
    latest_response = survey['responses'][-1]
    responses = latest_response['responses']
    
    context = f"""
IT Infrastructure Site Survey Results:

Client: {survey.get('client_name', 'Unknown')}
Site: {survey.get('site_name', 'Unknown')}

Survey Responses:
"""
    
    for q in survey['questions']:
        qid = q['id']
        answer = responses.get(qid, 'Not answered')
        context += f"\n{q['category']} - {q['question']}\nAnswer: {answer}\n"
    
    prompt = f"""You are a senior network architect and IT infrastructure consultant. Based on the following site survey, provide a comprehensive network infrastructure recommendation.

{context}

Generate a detailed IT infrastructure recommendation report including:

1. **EXECUTIVE SUMMARY**
   - Overview of the site and requirements
   - Recommended approach (high-level)

2. **NETWORK TOPOLOGY RECOMMENDATION**
   - Recommended network architecture (star, mesh, hybrid)
   - VLAN structure if applicable
   - Network segmentation strategy

3. **HARDWARE RECOMMENDATIONS**
   - Firewall/router specifications
   - Switch recommendations (managed/unmanaged, PoE needs)
   - Wireless access point placement and quantity
   - Cabling requirements

4. **INTERNET CONNECTIVITY**
   - Recommended ISP and bandwidth
   - Backup connectivity options
   - WAN configuration if multi-site

5. **SECURITY RECOMMENDATIONS**
   - Security appliances/services needed
   - Compliance considerations
   - Access control policies

6. **ESTIMATED COSTS**
   - Equipment costs
   - Installation costs
   - Ongoing service costs

7. **IMPLEMENTATION TIMELINE**
   - Phase 1: Planning and procurement
   - Phase 2: Installation
   - Phase 3: Testing and cutover

8. **PRIORITY ACTIONS**
   - Top 3 immediate actions needed

Format as a professional proposal suitable for client presentation.
"""

    def generate():
        try:
            resp = requests.post(
                f'{OLLAMA_BASE}/api/generate',
                json={'model': 'llama3.2:3b', 'prompt': prompt, 'stream': True},
                stream=True,
                timeout=180
            )
            for line in resp.iter_lines():
                if line:
                    yield line.decode('utf-8') + '\n'
        except Exception as e:
            error_chunk = json.dumps({'error': str(e), 'done': True})
            yield error_chunk + '\n'

    return Response(stream_with_context(generate()), mimetype='application/x-ndjson')


@app.route('/api/surveys/<survey_id>/generate-quote', methods=['POST'])
def generate_quote(survey_id):
    """Generate a service quote based on survey responses."""
    survey = surveys_db.get(survey_id)
    if not survey or not survey['responses']:
        return jsonify({'error': 'No survey data available'}), 400
    
    latest_response = survey['responses'][-1]
    responses = latest_response['responses']
    
    context = f"""
Site Survey Data for Quote Generation:

Client: {survey.get('client_name', 'Unknown')}
Site: {survey.get('site_name', 'Unknown')}
"""
    
    for q in survey['questions']:
        qid = q['id']
        answer = responses.get(qid, 'Not answered')
        context += f"\n{q['category']}: {answer}"
    
    prompt = f"""You are an IT services sales consultant. Based on the following site survey, create a professional service quote for network infrastructure setup.

{context}

Generate a professional service quote in the following format:

**IT INFRASTRUCTURE SERVICES QUOTE**
**For: {survey.get('client_name', 'Client')}**
**Site: {survey.get('site_name', 'Site')}**
**Date: {datetime.now().strftime('%B %d, %Y')}**

**SCOPE OF WORK:**
[List the specific services needed based on the survey - network design, equipment procurement, installation, configuration, testing]

**RECOMMENDED SOLUTION:**
1. [Category - e.g., Network Design & Planning] - $[Price]
   - [Detailed description of what's included]
2. [Category - e.g., Equipment & Hardware] - $[Price]
   - [List of equipment and costs]
3. [Category - e.g., Installation & Configuration] - $[Price]
   - [Labor hours and rates]
4. [Category - e.g., Security Implementation] - $[Price]
   - [Security services included]
5. [Category - e.g., Testing & Documentation] - $[Price]
   - [Testing procedures and documentation]

**INVESTMENT SUMMARY:**
- Initial Assessment & Design: $[Amount]
- Equipment/Hardware: $[Amount]
- Labor/Installation: $[Amount]
- Training & Documentation: $[Amount]
- **Total Project Investment: $[Amount]**

**OPTIONAL ONGOING SERVICES:**
- Monthly Management & Monitoring: $[Amount]/month
- Annual Maintenance: $[Amount]/year
- Support Contract: $[Amount]/month

**TIMELINE:** [Timeframe]
- Week 1-2: Design & Procurement
- Week 3-4: Installation
- Week 5: Testing & Training

**TERMS:**
- 50% deposit required to begin
- Net 15 on completion
- Quote valid for 30 days

This is an estimate based on the initial site survey. Final pricing will be confirmed after detailed site assessment.
"""

    response = call_ollama(prompt, stream=False)
    return jsonify({'quote': response})


@app.route('/api/surveys/<survey_id>/export/pdf', methods=['POST'])
def export_pdf(survey_id):
    """Export survey report to PDF."""
    survey = surveys_db.get(survey_id)
    if not survey or not survey['responses']:
        return jsonify({'error': 'No survey data available'}), 400
    
    latest_response = survey['responses'][-1]
    responses = latest_response['responses']
    
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=letter,
                          rightMargin=72, leftMargin=72,
                          topMargin=72, bottomMargin=18)
    
    elements = []
    styles = getSampleStyleSheet()
    
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        textColor=colors.HexColor('#1f6feb'),
        spaceAfter=30
    )
    
    heading_style = ParagraphStyle(
        'CustomHeading',
        parent=styles['Heading2'],
        fontSize=16,
        textColor=colors.HexColor('#58a6ff'),
        spaceAfter=12,
        spaceBefore=12
    )
    
    # Title
    elements.append(Paragraph(f"IT Site Survey Report", title_style))
    elements.append(Paragraph(f"Client: {survey.get('client_name', 'Unknown')}", styles['Normal']))
    elements.append(Paragraph(f"Site: {survey.get('site_name', 'Unknown')}", styles['Normal']))
    elements.append(Paragraph(f"Date: {datetime.now().strftime('%B %d, %Y')}", styles['Normal']))
    elements.append(Spacer(1, 0.3*inch))
    
    # Survey Responses
    elements.append(Paragraph("Survey Responses", heading_style))
    
    for q in survey['questions']:
        qid = q['id']
        answer = responses.get(qid, 'Not answered')
        if isinstance(answer, list):
            answer = ', '.join(answer)
        
        elements.append(Paragraph(f"<b>{q['category']}</b>", styles['Heading3']))
        elements.append(Paragraph(q['question'], styles['Normal']))
        elements.append(Paragraph(f"Answer: {answer}", styles['Normal']))
        elements.append(Spacer(1, 0.1*inch))
    
    doc.build(elements)
    
    buffer.seek(0)
    return send_file(
        buffer,
        as_attachment=True,
        download_name=f'Site_Survey_{survey.get("client_name", "Client").replace(" ", "_")}_{datetime.now().strftime("%Y%m%d")}.pdf',
        mimetype='application/pdf'
    )


@app.route('/api/surveys/<survey_id>/export/text', methods=['POST'])
def export_text(survey_id):
    """Export survey report to formatted text."""
    survey = surveys_db.get(survey_id)
    if not survey or not survey['responses']:
        return jsonify({'error': 'No survey data available'}), 400
    
    latest_response = survey['responses'][-1]
    responses = latest_response['responses']
    
    text_report = f"""IT SITE SURVEY REPORT
{'=' * 60}

Client: {survey.get('client_name', 'Unknown')}
Site: {survey.get('site_name', 'Unknown')}
Date: {datetime.now().strftime('%B %d, %Y')}

{'=' * 60}
SURVEY RESPONSES:
{'=' * 60}

"""
    
    for q in survey['questions']:
        qid = q['id']
        answer = responses.get(qid, 'Not answered')
        if isinstance(answer, list):
            answer = ', '.join(answer)
        
        text_report += f"""
{q['category']}:
Question: {q['question']}
Answer: {answer}

"""
    
    text_report += f"""
{'=' * 60}
Report generated by IT Site Survey AI
{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    buffer = io.BytesIO(text_report.encode('utf-8'))
    return send_file(
        buffer,
        as_attachment=True,
        download_name=f'Site_Survey_{survey.get("client_name", "Client").replace(" ", "_")}_{datetime.now().strftime("%Y%m%d")}.txt',
        mimetype='text/plain'
    )


@app.route('/api/surveys/<survey_id>/export/email', methods=['POST'])
def export_email(survey_id):
    """Export survey report formatted for email."""
    survey = surveys_db.get(survey_id)
    if not survey or not survey['responses']:
        return jsonify({'error': 'No survey data available'}), 400
    
    latest_response = survey['responses'][-1]
    responses = latest_response['responses']
    
    email_body = f"""
<h2>IT Site Survey Report - {survey.get('client_name', 'Client')}</h2>
<p><strong>Site:</strong> {survey.get('site_name', 'Unknown')}</p>
<p><strong>Date:</strong> {datetime.now().strftime('%B %d, %Y')}</p>

<hr>

<h3>Survey Responses:</h3>

<table border="1" cellpadding="5" cellspacing="0" style="border-collapse: collapse;">
  <tr style="background-color: #f0f0f0;">
    <th>Category</th>
    <th>Question</th>
    <th>Answer</th>
  </tr>
"""
    
    for q in survey['questions']:
        qid = q['id']
        answer = responses.get(qid, 'Not answered')
        if isinstance(answer, list):
            answer = ', '.join(answer)
        
        email_body += f"""
  <tr>
    <td>{q['category']}</td>
    <td>{q['question']}</td>
    <td>{answer}</td>
  </tr>
"""
    
    email_body += """
</table>

<p>---<br>
Report generated by IT Site Survey AI<br>
""" + datetime.now().strftime('%Y-%m-%d %H:%M:%S') + "</p>
"""
    
    return jsonify({'email_body': email_body})


if __name__ == '__main__':
    port = int(os.environ.get('PORT', 3003))
    print("───────────────────────────────────────")
    print(" 🏢 IT Site Survey AI")
    print(f" Running on http://0.0.0.0:{port}")
    print(f" Ollama: {OLLAMA_BASE}")
    print("───────────────────────────────────────")
    app.run(host='0.0.0.0', port=port, debug=False, use_reloader=False)