#!/usr/bin/env python3
"""
IT Site Survey AI - Flask Backend API
"""

import os
import json
import uuid
import time
from datetime import datetime
from flask import Flask, request, jsonify, Response, render_template, send_from_directory
from flask_cors import CORS
import base64
import requests
import io
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.lib import colors
from reportlab.lib.enums import TA_LEFT, TA_CENTER
import urllib.parse

app = Flask(__name__)
CORS(app)

# In-memory storage
surveys = {}

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

def get_ollama_models():
    """Fetch available models from Ollama"""
    try:
        resp = requests.get(f'{OLLAMA_BASE}/api/tags', timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            models = []
            for model in data.get('models', []):
                model_name = model.get('name', model.get('model', ''))
                models.append({
                    "id": model_name,
                    "name": model_name,
                    "provider": "Ollama"
                })
            return models
    except Exception as e:
        print(f"Error fetching Ollama models: {e}")
    return [{"id": "llama3.2:latest", "name": "Llama 3.2", "provider": "Ollama"}]

MODELS = get_ollama_models()

# 35-question survey template
SURVEY_TEMPLATE = {
    "metadata": {
        "title": "Network Design and Installation Site Survey",
        "version": "3.0",
        "total_questions": 42,
        "source": "Network Cabling Site Survey Template PDF"
    },
    "questions": [
        {
            "id": "q1",
            "category": "Client Information",
            "question": "Company Name",
            "type": "text"
        },
        {
            "id": "q2",
            "category": "Client Information",
            "question": "Industry/Vertical",
            "type": "select",
            "options": [
                "Manufacturing",
                "Healthcare",
                "Education (K-12)",
                "Education (Higher Ed)",
                "Financial Services",
                "Retail",
                "Warehouse/Distribution",
                "Technology/SaaS",
                "Government",
                "Legal",
                "Real Estate",
                "Hospitality",
                "Transportation/Logistics",
                "Energy/Utilities",
                "Construction",
                "Non-profit",
                "Professional Services",
                "Media/Entertainment",
                "Agriculture",
                "Other"
            ]
        },
        {
            "id": "q3",
            "category": "Client Information",
            "question": "Primary Business Function",
            "type": "select",
            "options": [
                "Office/Administrative",
                "Manufacturing/Production",
                "Research & Development",
                "Data Center/Colocation",
                "Call Center",
                "Warehouse/Storage",
                "Retail Storefront",
                "Mixed-Use Facility",
                "Medical/Patient Care",
                "Laboratory",
                "Clean Room",
                "Food Processing",
                "Assembly Line",
                "Distribution Center",
                "Corporate Headquarters",
                "Training/Education Facility",
                "Other"
            ]
        },
        {
            "id": "q4",
            "category": "Client Information",
            "question": "Number of Employees",
            "type": "select",
            "options": [
                "1-10",
                "11-25",
                "26-50",
                "51-100",
                "101-250",
                "251-500",
                "501-1000",
                "1001-2500",
                "2501-5000",
                "5001-10000",
                "10000+"
            ]
        },
        {
            "id": "q5",
            "category": "Site Details",
            "question": "Site Name/Location",
            "type": "text"
        },
        {
            "id": "q6",
            "category": "Site Details",
            "question": "Site Address",
            "type": "text"
        },
        {
            "id": "q7",
            "category": "Site Details",
            "question": "Building Type",
            "type": "select",
            "options": [
                "Single Office",
                "Multi-tenant",
                "Campus",
                "Warehouse",
                "Retail",
                "Industrial",
                "Data Center",
                "Other"
            ]
        },
        {
            "id": "q8",
            "category": "Site Details",
            "question": "Total Square Footage",
            "type": "select",
            "options": [
                "Under 1,000",
                "1,000 - 5,000",
                "5,001 - 10,000",
                "10,001 - 25,000",
                "25,001 - 50,000",
                "50,001 - 100,000",
                "100,001 - 250,000",
                "250,001 - 500,000",
                "500,000+"
            ]
        },
        {
            "id": "q9",
            "category": "Site Details",
            "question": "Number of Buildings",
            "type": "select",
            "options": [
                "1",
                "2",
                "3-5",
                "6-10",
                "11-25",
                "26-50",
                "50+"
            ]
        },
        {
            "id": "q10",
            "category": "Site Details",
            "question": "Number of Floors per Building",
            "type": "select",
            "options": [
                "1 (Single-story)",
                "2",
                "3",
                "4-5",
                "6-10",
                "11-20",
                "21-50",
                "50+"
            ]
        },
        {
            "id": "q11",
            "category": "Business Operations",
            "question": "Operating Hours",
            "type": "text"
        },
        {
            "id": "q12",
            "category": "Business Operations",
            "question": "Shift Operations",
            "type": "select",
            "options": [
                "Single Shift",
                "Multiple Shifts",
                "24/7 Operations"
            ]
        },
        {
            "id": "q13",
            "category": "Business Operations",
            "question": "Acceptable Downtime Windows",
            "type": "text"
        },
        {
            "id": "q14",
            "category": "Business Operations",
            "question": "Business Continuity Requirements",
            "type": "select",
            "options": [
                "Standard",
                "High Availability",
                "Mission Critical"
            ]
        },
        {
            "id": "q15",
            "category": "Special Requirements",
            "question": "Special Requirements and Constraints (select all that apply)",
            "type": "multiselect",
            "options": [
                "Historic building restrictions",
                "Hazardous/explosive environment",
                "Clean room requirements",
                "Food processing facility",
                "Healthcare/patient care areas",
                "Educational institution",
                "Government facility/clearance required",
                "Multi-tenant with shared infrastructure",
                "Aesthetic restrictions (visible equipment)",
                "Other"
            ]
        },
        {
            "id": "q16",
            "category": "Existing Network Infrastructure",
            "question": "Current Network Topology",
            "type": "select",
            "options": [
                "None",
                "Flat",
                "Two-tier",
                "Three-tier",
                "Spine-Leaf",
                "Other"
            ]
        },
        {
            "id": "q17",
            "category": "Existing Network Infrastructure",
            "question": "Network Size",
            "type": "select",
            "options": [
                "Small (<50 devices)",
                "Medium (50-250)",
                "Large (250-1000)",
                "Enterprise (1000+)"
            ]
        },
        {
            "id": "q18",
            "category": "Existing Network Infrastructure",
            "question": "Internet Service Provider(s)",
            "type": "select",
            "options": [
                "Comcast Business",
                "AT&T",
                "Verizon",
                "Spectrum",
                "Cox Business",
                "Frontier",
                "CenturyLink/Lumen",
                "Charter Business",
                "Windstream",
                "Google Fiber",
                "Other"
            ]
        },
        {
            "id": "q19",
            "category": "Existing Network Infrastructure",
            "question": "Internet Bandwidth (Primary and Backup)",
            "type": "select",
            "options": [
                "10 Mbps",
                "25 Mbps",
                "50 Mbps",
                "100 Mbps",
                "250 Mbps",
                "500 Mbps",
                "1 Gbps",
                "2.5 Gbps",
                "5 Gbps",
                "10 Gbps",
                "25 Gbps",
                "40 Gbps",
                "100 Gbps",
                "Need Assessment"
            ]
        },
        {
            "id": "q20",
            "category": "Existing Network Infrastructure",
            "question": "ISP Connection Type",
            "type": "select",
            "options": [
                "Fiber",
                "Cable",
                "DSL",
                "T1/T3",
                "5G/LTE",
                "Satellite"
            ]
        },
        {
            "id": "q21",
            "category": "Existing Equipment",
            "question": "Active Network Equipment Inventory (list make/model, quantity, location, age, condition)",
            "type": "textarea"
        },
        {
            "id": "q22",
            "category": "Existing Equipment",
            "question": "Total Active Ports",
            "type": "select",
            "options": [
                "1-12",
                "13-24",
                "25-36",
                "37-48"
            ]
        },
        {
            "id": "q22b",
            "category": "Existing Equipment",
            "question": "Unused/Available Ports",
            "type": "select",
            "options": [
                "1-12",
                "13-24",
                "25-36",
                "37-48"
            ]
        },
        {
            "id": "q23",
            "category": "Existing Equipment",
            "question": "Port Types and Speeds",
            "type": "multiselect",
            "options": [
                "RJ45",
                "Fiber",
                "100 Mbps",
                "1 Gbps",
                "10 Gbps",
                "Multi-Gig",
                "Mixed"
            ]
        },
        {
            "id": "q24",
            "category": "Telecommunications Room",
            "question": "Main Distribution Frame (MDF) Location and Details",
            "type": "textarea"
        },
        {
            "id": "q25",
            "category": "Telecommunications Room",
            "question": "Intermediate Distribution Frames (IDFs) - Number and Locations",
            "type": "textarea"
        },
        {
            "id": "q26",
            "category": "Server Infrastructure",
            "question": "Server Location",
            "type": "select",
            "options": [
                "On-premises",
                "Cloud",
                "Hybrid",
                "Colocation"
            ]
        },
        {
            "id": "q27",
            "category": "Server Infrastructure",
            "question": "Server Room/Data Center Details (size, cooling, power, fire suppression)",
            "type": "textarea"
        },
        {
            "id": "q28",
            "category": "Wireless Infrastructure",
            "question": "Existing Wireless Infrastructure",
            "type": "select",
            "options": [
                "None",
                "Consumer Grade",
                "Enterprise",
                "Needs Assessment"
            ]
        },
        {
            "id": "q29",
            "category": "Wireless Infrastructure",
            "question": "Wireless Standard",
            "type": "select",
            "options": [
                "802.11n (Wi-Fi 4)",
                "802.11ac (Wi-Fi 5)",
                "802.11ax (Wi-Fi 6)",
                "802.11be (Wi-Fi 7)",
                "Mixed Standards"
            ]
        },
        {
            "id": "q29b",
            "category": "Wireless Infrastructure",
            "question": "Number of Access Points",
            "type": "select",
            "options": [
                "1-10",
                "11-25",
                "26-50",
                "51-100",
                "101-250",
                "250+"
            ]
        },
        {
            "id": "q30",
            "category": "Wireless Infrastructure",
            "question": "Wireless Coverage Issues",
            "type": "multiselect",
            "options": [
                "Dead zones",
                "Interference",
                "Insufficient capacity",
                "No issues"
            ]
        },
        {
            "id": "q31",
            "category": "Voice Systems",
            "question": "Phone System Type",
            "type": "select",
            "options": [
                "VoIP",
                "Traditional PBX",
                "Cloud/Hosted",
                "None"
            ]
        },
        {
            "id": "q31b",
            "category": "Voice Systems",
            "question": "Number of Phone Extensions",
            "type": "select",
            "options": [
                "1-10",
                "11-25",
                "26-50",
                "51-100",
                "101-250",
                "251-500",
                "500+"
            ]
        },
        {
            "id": "q32",
            "category": "Security Systems",
            "question": "Physical Security Systems",
            "type": "multiselect",
            "options": [
                "IP Camera System",
                "Access Control System (Badge/Card)",
                "Biometric Access Control",
                "Intrusion Detection/Alarm",
                "Building Management System (BMS)",
                "Intercom/Paging System",
                "Video Door Entry",
                "Motion Sensors",
                "Glass Break Detectors",
                "Panic Buttons",
                "Security Gates/Barriers",
                "None"
            ]
        },
        {
            "id": "q33",
            "category": "Network Requirements",
            "question": "Primary Project Objectives (up to 3)",
            "type": "textarea"
        },
        {
            "id": "q34",
            "category": "Network Requirements",
            "question": "Required Network Speed",
            "type": "select",
            "options": [
                "1 Gbps",
                "10 Gbps",
                "25 Gbps",
                "40 Gbps",
                "100 Gbps",
                "Multi-Gig (2.5/5 Gbps)"
            ]
        },
        {
            "id": "q35",
            "category": "Network Requirements",
            "question": "Wireless Requirements",
            "type": "multiselect",
            "options": [
                "Guest Access",
                "BYOD Support",
                "IoT Devices",
                "High Density",
                "Outdoor Coverage"
            ]
        },
        {
            "id": "q36",
            "category": "Network Requirements",
            "question": "Uptime Requirement",
            "type": "select",
            "options": [
                "99%",
                "99.9%",
                "99.99%",
                "99.999%"
            ]
        },
        {
            "id": "q37",
            "category": "Network Requirements",
            "question": "Redundancy Level",
            "type": "select",
            "options": [
                "None",
                "Basic (Dual ISP)",
                "High (Redundant Paths)",
                "Full (N+1 Equipment)"
            ]
        },
        {
            "id": "q38",
            "category": "Applications",
            "question": "Critical Business Applications (select all that apply)",
            "type": "multiselect",
            "options": [
                "Email/Microsoft 365/Google Workspace",
                "ERP System (SAP, Oracle, NetSuite)",
                "CRM System (Salesforce, HubSpot, Dynamics)",
                "Database Applications (SQL Server, MySQL, PostgreSQL)",
                "File Sharing/Collaboration (SharePoint, Dropbox, Box)",
                "Cloud Applications (AWS, Azure, GCP)",
                "Video Conferencing (Zoom, Teams, Webex, Google Meet)",
                "VoIP Telephony (RingCentral, 8x8, Vonage)",
                "Remote Desktop/VDI (Citrix, VMware Horizon, RDS)",
                "Video Surveillance (Genetec, Milestone, Avigilon)",
                "Point of Sale (POS) Systems",
                "Industrial Control Systems/SCADA",
                "Media Streaming/Content Delivery",
                "Backup and Disaster Recovery (Veeam, Commvault, Rubrik)"
            ]
        },
        {
            "id": "q39",
            "category": "Device Count",
            "question": "Device Count and Network Endpoints (list current count by category)",
            "type": "textarea"
        },
        {
            "id": "q40",
            "category": "Wireless Design",
            "question": "Wireless Coverage Areas and Requirements (sq ft, user density, client types)",
            "type": "textarea"
        },
        {
            "id": "q41",
            "category": "Security",
            "question": "Security Priorities (threat protection, access control, VPN requirements)",
            "type": "textarea"
        },
        {
            "id": "q42",
            "category": "Compliance",
            "question": "Compliance Requirements (select all that apply)",
            "type": "multiselect",
            "options": [
                "HIPAA (Healthcare)",
                "PCI-DSS (Payment Card Industry)",
                "SOX (Sarbanes-Oxley)",
                "FERPA (Educational)",
                "NIST/CMMC (Defense Contractor)",
                "GDPR (European Privacy)",
                "CCPA (California Privacy)",
                "ISO 27001",
                "SOC 2",
                "FISMA (Federal)",
                "GLBA (Financial Services)",
                "Industry-specific compliance",
                "None"
            ]
        }
    ]
}
# Sample data for testing
SAMPLE_DATA = {
    "name": "Network Design Site Survey - ABC Manufacturing",
    "client": "ABC Manufacturing Inc.",
    "site": "Main Plant - Detroit, MI",
    "description": "Network infrastructure upgrade for manufacturing facility expansion",
    "model": "llama3.2:3b",
    "answers": {
        "q1": "ABC Manufacturing Inc.",
        "q2": "Sample data for Industry/Vertical...",
        "q3": "Sample data for Primary Business Function...",
        "q4": "350 employees",
        "q5": "ABC Manufacturing Inc.",
        "q6": "1234 Industrial Blvd, Detroit, MI 48201",
        "q7": "Industrial",
        "q8": "125,000 sq ft",
        "q9": "2 (Main Plant + Warehouse)",
        "q10": "3 floors in main plant, 1 in warehouse",
        "q11": "6:00 AM - 11:00 PM Mon-Sat",
        "q12": "Multiple Shifts",
        "q13": "Sunday 12:00 AM - 6:00 AM",
        "q14": "High Availability",
        "q15": [
            "Industrial",
            "Food processing facility"
        ],
        "q16": "Three-tier",
        "q17": "Medium (50-250)",
        "q18": "Sample data for Internet Service Provider(s)...",
        "q19": "Primary: 1 Gbps, Backup: 500 Mbps",
        "q20": "Fiber",
        "q21": "Core Switches: Cisco C9300-48P (2) - 3 years old, good condition\\nAccess Switches: Cisco C9200-24P (12) - 2 years old, good condition\\nFirewall: Fortinet FortiGate 600E (1) - 1 year old, excellent condition\\nRouter: Cisco ISR 4331 (1) - 4 years old, fair condition",
        "q22": "Total: 288 ports, Available: 47 ports",
        "q23": [
            "RJ45",
            "Fiber",
            "1 Gbps",
            "10 Gbps"
        ],
        "q24": "MDF: Building A Room 101\\n- Size: 12x16 ft\\n- Available Rack Space: 30U of 42U\\n- Power: 120V, 208V, UPS (10kVA), Generator backup\\n- Cooling: Adequate (dedicated HVAC)\\n- Security: Locked with badge access\\n\\nIDF: Building B Room 205\\n- Size: 8x10 ft\\n- Available Rack Space: 8U of 24U\\n- Power: 120V, UPS (3kVA)\\n- Cooling: Portable AC unit (needs upgrade)",
        "q25": "MDF: Building A Room 101\\n- Size: 12x16 ft\\n- Available Rack Space: 30U of 42U\\n- Power: 120V, 208V, UPS (10kVA), Generator backup\\n- Cooling: Adequate (dedicated HVAC)\\n- Security: Locked with badge access\\n\\nIDF: Building B Room 205\\n- Size: 8x10 ft\\n- Available Rack Space: 8U of 24U\\n- Power: 120V, UPS (3kVA)\\n- Cooling: Portable AC unit (needs upgrade)",
        "q26": "On-premises",
        "q27": "Server Room: 200 sq ft\\n- Available Rack Space: 15U of 24U\\n- Cooling: Dedicated split AC unit, 36,000 BTU\\n- Power: UPS (20kVA), Dual feeds, Generator backup\\n- Fire Suppression: FM-200 system installed",
        "q28": "Enterprise",
        "q29": "802.11ac (Wi-Fi 5), 24 APs total",
        "q30": [
            "Dead zones",
            "Interference"
        ],
        "q31": "VoIP system, 275 extensions, Cisco CUCM",
        "q32": [
            "IP Camera System",
            "Access Control System",
            "Intrusion Detection/Alarm"
        ],
        "q33": "1. Upgrade entire facility to Wi-Fi 6 (802.11ax) with seamless roaming\\n2. Implement network segmentation to isolate OT networks from IT networks for security\\n3. Deploy redundant internet connectivity with SD-WAN for improved reliability",
        "q34": "10 Gbps",
        "q35": [
            "BYOD Support",
            "IoT Devices",
            "High Density"
        ],
        "q36": "99.99%",
        "q37": "High (Redundant Paths)",
        "q38": [
            "Email/Microsoft 365/Google Workspace",
            "ERP System",
            "File sharing/collaboration",
            "Video conferencing (Zoom, Teams, Webex)",
            "VoIP telephony"
        ],
        "q39": "Desktop Computers: 180\\nLaptops: 95\\nTablets/Mobile: 45\\nVoIP Phones: 275\\nPrinters/MFPs: 45\\nWireless APs: 24\\nIP Cameras: 64\\nAccess Control: 32\\nIoT Devices: 120 (temperature sensors, HVAC controls)\\nManufacturing/Industrial: 85 (PLC controllers, HMI panels)\\nPoint of Sale: 12\\nGuest Devices: Variable (20-50)",
        "q40": "Office Space: 45,000 sq ft - Medium density (15-20 devices per AP), laptops, phones, tablets\\nWarehouse/Production: 80,000 sq ft - Low density, high ceilings (25 ft), handheld scanners, IoT sensors\\nOutdoor Areas: 5,000 sq ft - Loading docks, basic coverage for mobile devices\\nUser Density: Medium (10-30 devices per AP) in office, Low in warehouse",
        "q41": "Threat Protection: Next-Gen Firewall (NGFW) with IPS\\nNetwork Access Control: 802.1X with NAC solution\\nVPN Requirements: Site-to-Site VPN for 3 remote offices, Remote Access VPN for 50 users\\nNetwork Monitoring: Full NPM solution with SNMP, NetFlow, and SIEM integration",
        "q42": [
            "HIPAA (Healthcare)",
            "PCI-DSS (Payment Card)"
        ]
    }
}

# Route for serving the frontend
@app.route('/')
def index():
    """Serve the main HTML page"""
    return send_from_directory('templates', 'index.html')

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

# API Routes
@app.route('/api/models', methods=['GET'])
def get_models():
    """Get available AI models"""
    models = get_ollama_models()
    return jsonify({"models": models, "default": models[0]["id"] if models else "llama3.2:latest"})

@app.route('/api/surveys/template', methods=['GET'])
def get_template():
    """Get survey template with all 35 questions"""
    return jsonify(SURVEY_TEMPLATE)

@app.route('/api/surveys', methods=['GET'])
def get_surveys():
    """Get all surveys"""
    survey_list = []
    for survey_id, survey in surveys.items():
        survey_list.append({
            "id": survey_id,
            "name": survey.get("name", "Unnamed Survey"),
            "client": survey.get("client", ""),
            "site": survey.get("site", ""),
            "created_at": survey.get("created_at", ""),
            "updated_at": survey.get("updated_at", ""),
            "has_analysis": bool(survey.get("analysis", ""))
        })
    # Sort by created_at descending
    survey_list.sort(key=lambda x: x.get("created_at", ""), reverse=True)
    return jsonify({"surveys": survey_list})

@app.route('/api/surveys', methods=['POST'])
def create_survey():
    """Create a new survey"""
    data = request.get_json()
    
    survey_id = str(uuid.uuid4())
    survey = {
        "id": survey_id,
        "name": data.get("name", ""),
        "client": data.get("client", ""),
        "site": data.get("site", ""),
        "description": data.get("description", ""),
        "model": data.get("model", "gpt-4o"),
        "answers": data.get("answers", {}),
        "photos": data.get("photos", []),
        "created_at": datetime.now().isoformat(),
        "updated_at": datetime.now().isoformat()
    }
    
    surveys[survey_id] = survey
    
    return jsonify({"id": survey_id, "message": "Survey created successfully"}), 201

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

@app.route('/api/surveys/<survey_id>', methods=['PUT'])
def update_survey(survey_id):
    """Update a survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json()
    survey = surveys[survey_id]
    
    survey["name"] = data.get("name", survey.get("name", ""))
    survey["client"] = data.get("client", survey.get("client", ""))
    survey["site"] = data.get("site", survey.get("site", ""))
    survey["description"] = data.get("description", survey.get("description", ""))
    survey["model"] = data.get("model", survey.get("model", "gpt-4o"))
    survey["answers"] = data.get("answers", survey.get("answers", {}))
    survey["photos"] = data.get("photos", survey.get("photos", []))
    survey["updated_at"] = datetime.now().isoformat()
    
    return jsonify({"message": "Survey updated successfully"})

@app.route('/api/surveys/<survey_id>', methods=['DELETE'])
def delete_survey(survey_id):
    """Delete a survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    del surveys[survey_id]
    return jsonify({"message": "Survey deleted successfully"})

@app.route('/api/surveys', methods=['DELETE'])
def delete_all_surveys():
    """Delete all surveys"""
    surveys.clear()
    return jsonify({"message": "All surveys deleted successfully"})

@app.route('/api/surveys/<survey_id>/analyze', methods=['POST'])
def analyze_survey(survey_id):
    """Generate AI analysis for a survey (streaming)"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json() or {}
    model = data.get("model", surveys[survey_id].get("model", "llama3.2:latest"))
    
    survey = surveys[survey_id]
    
    # Build survey context
    context = f"""
# IT Site Survey Analysis

**Organization:** {survey.get('client', 'N/A')}
**Site:** {survey.get('site', 'N/A')}
**Survey:** {survey.get('name', 'N/A')}

## Survey Responses Summary:
"""
    for q in SURVEY_TEMPLATE["questions"]:
        qid = q["id"]
        answer = survey.get("answers", {}).get(qid, "Not answered")
        context += f"\n**Q{qid}: {q['question']}**\nA: {answer}\n"
    
    prompt = f"""You are a network infrastructure expert. Based on the following site survey data, provide a comprehensive analysis with actionable recommendations.

{context}

Please provide:
1. Executive Summary
2. Critical Findings
3. Immediate Actions (0-3 months)
4. Short-term Initiatives (3-6 months)
5. Strategic Roadmap (6-18 months)
6. Resource Planning
7. Compliance Considerations

Be specific and reference actual survey responses in your analysis."""

    def generate():
        try:
            resp = requests.post(
                f'{OLLAMA_BASE}/api/generate',
                json={'model': model, 'prompt': prompt, 'stream': True},
                stream=True,
                timeout=180
            )
            for line in resp.iter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        if 'response' in data:
                            yield data['response']
                        if data.get('done'):
                            break
                    except:
                        pass
        except Exception as e:
            yield f"Error generating analysis: {str(e)}"
    
    return Response(generate(), mimetype='text/plain')

@app.route('/api/surveys/<survey_id>/chat', methods=['POST'])
def chat_about_survey(survey_id):
    """Chat with AI about a survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json()
    message = data.get("message", "")
    model = data.get("model", surveys[survey_id].get("model", "gpt-4o"))
    
    survey = surveys[survey_id]
    
    # Simple keyword-based responses for demo
    response = ""
    msg_lower = message.lower()
    
    if "budget" in msg_lower or "cost" in msg_lower or "spend" in msg_lower:
        response = f"Based on the survey, the current monthly cloud spend is approximately $45,000. With optimization efforts already underway using Reserved Instances, there's potential for 20-30% cost savings. The organization should also factor in costs for the planned SD-WAN deployment and cloud migration expansion."
    elif "security" in msg_lower or "threat" in msg_lower or "risk" in msg_lower:
        response = f"The security assessment shows strong fundamentals with Palo Alto Networks firewalls and CrowdStrike endpoint protection. Key gaps include: MFA not fully deployed across all systems, incident response plan not yet tested, and the planned Zero Trust architecture still in planning. Recommend prioritizing MFA completion and scheduling an incident response tabletop exercise."
    elif "cloud" in msg_lower or "aws" in msg_lower or "azure" in msg_lower:
        response = f"Currently at 30% cloud adoption with AWS as primary and Azure for Microsoft workloads. The target is 80% cloud migration by end of 2025. Key next steps include completing the cloud governance framework and optimizing the current $45K/month spend before expanding workloads."
    elif "network" in msg_lower or "bandwidth" in msg_lower or "connectivity" in msg_lower:
        response = f"Network infrastructure includes 1Gbps at HQ, 500Mbps at Denver, and 300Mbps at Austin. SD-WAN is in planning stages, which will improve inter-site connectivity and provide redundancy. Critical: Address single points of failure at Denver (single internet circuit) and Austin (power redundancy)."
    elif "legacy" in msg_lower or "as400" in msg_lower or "modernization" in msg_lower:
        response = f"The AS/400 inventory system is identified as a legacy application requiring modernization. Recommended approach: 1) Implement an API layer as short-term bridge, 2) Evaluate cloud-based ERP alternatives, 3) Plan phased migration to avoid disruption. This should be included in the FY initiatives."
    elif "staff" in msg_lower or "team" in msg_lower or "resource" in msg_lower:
        response = f"The current IT staffing ratio is 1:75 (6 staff for 450 users), which is lean. Average ticket resolution times are good at 4.2 hours for L2. For the ambitious digital transformation initiatives planned, consider adding 2-3 FTEs focused on cloud and security, or engage managed security services for 24/7 coverage."
    elif "backup" in msg_lower or "disaster" in msg_lower or "recovery" in msg_lower:
        response = f"Backup solution is Veeam with RTO of 4 hours and RPO of 1 hour. This meets most business requirements. Recommend validating these metrics with annual DR testing and ensure cloud workloads have equivalent protection through AWS/Azure backup services."
    elif "compliance" in msg_lower or "audit" in msg_lower or "soc2" in msg_lower or "iso" in msg_lower:
        response = f"The organization maintains SOC2 Type II, ISO 27001, and PCI-DSS compliance. The last SOC2 audit was 2 months ago with minor findings already remediated. Data classification policy is implemented. Ensure all planned infrastructure changes maintain these compliance standards."
    else:
        response = f"I've reviewed the survey for {survey.get('client', 'your organization')}. With 450 employees across 3 locations, you're managing a hybrid infrastructure (70% on-prem, 30% cloud) with ambitious goals including 80% cloud migration and Zero Trust implementation. Your security stack is enterprise-grade but MFA needs to be expanded. The 1:75 IT staff ratio is lean for these initiatives - consider additional resources or managed services. What specific aspect would you like to explore further?"
    
    return jsonify({
        "response": response,
        "model": model
    })

@app.route('/api/surveys/<survey_id>/analysis', methods=['POST'])
def save_analysis(survey_id):
    """Save AI analysis to survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json() or {}
    analysis_content = data.get("analysis", "")
    
    surveys[survey_id]["analysis"] = analysis_content
    surveys[survey_id]["updated_at"] = datetime.now().isoformat()
    
    return jsonify({"message": "Analysis saved successfully"})

@app.route('/api/surveys/<survey_id>/analysis', methods=['DELETE'])
def delete_analysis(survey_id):
    """Delete AI analysis from survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    if "analysis" in surveys[survey_id]:
        del surveys[survey_id]["analysis"]
        surveys[survey_id]["updated_at"] = datetime.now().isoformat()
        return jsonify({"message": "Analysis deleted successfully"})
    
    return jsonify({"error": "No analysis found"}), 404

@app.route('/api/surveys/<survey_id>/analysis', methods=['GET'])
def get_analysis(survey_id):
    """Get saved AI analysis for a survey"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    analysis = surveys[survey_id].get("analysis", "")
    return jsonify({"analysis": analysis, "has_analysis": bool(analysis)})

@app.route('/api/surveys/<survey_id>/export', methods=['POST'])
def export_survey(survey_id):
    """Export survey in various formats"""
    if survey_id not in surveys:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json() or {}
    format_type = data.get("format", "text")
    content = data.get("content", "")
    
    survey = surveys[survey_id]
    
    # If no content provided, check if survey has saved analysis
    if not content and survey.get("analysis"):
        content = survey["analysis"]
    
    if format_type == "text":
        # Plain text export
        export_content = f"""
IT SITE SURVEY REPORT
{'=' * 50}

Survey Name: {survey.get('name', 'N/A')}
Client: {survey.get('client', 'N/A')}
Site: {survey.get('site', 'N/A')}
Description: {survey.get('description', 'N/A')}
Created: {survey.get('created_at', 'N/A')}

{'=' * 50}
SURVEY RESPONSES
{'=' * 50}
"""
        for q in SURVEY_TEMPLATE["questions"]:
            qid = q["id"]
            answer = survey.get("answers", {}).get(qid, "Not answered")
            export_content += f"\n\nQ: {q['question']}\nA: {answer}"
        
        if content:
            export_content += f"\n\n{'=' * 50}\nAI ANALYSIS\n{'=' * 50}\n\n{content}"
        
        return jsonify({
            "format": "text",
            "filename": f"{survey.get('name', 'survey').replace(' ', '_')}_report.txt",
            "content": export_content
        })
    
    elif format_type == "json":
        # JSON export
        export_data = {
            "survey": survey,
            "template": SURVEY_TEMPLATE,
            "analysis": content
        }
        return jsonify({
            "format": "json",
            "filename": f"{survey.get('name', 'survey').replace(' ', '_')}_report.json",
            "content": json.dumps(export_data, indent=2)
        })
    
    elif format_type == "pdf":
        # Generate PDF using reportlab
        buffer = io.BytesIO()
        doc = SimpleDocTemplate(
            buffer,
            pagesize=letter,
            rightMargin=72,
            leftMargin=72,
            topMargin=72,
            bottomMargin=18
        )
        
        # Container for the 'Flowable' objects
        elements = []
        
        # Styles
        styles = getSampleStyleSheet()
        title_style = ParagraphStyle(
            'CustomTitle',
            parent=styles['Heading1'],
            fontSize=24,
            textColor=colors.HexColor('#1e40af'),
            spaceAfter=30,
            alignment=TA_CENTER
        )
        heading_style = ParagraphStyle(
            'CustomHeading',
            parent=styles['Heading2'],
            fontSize=14,
            textColor=colors.HexColor('#1e40af'),
            spaceAfter=12,
            spaceBefore=12
        )
        normal_style = styles["Normal"]
        normal_style.fontSize = 10
        normal_style.leading = 14
        
        # Title
        elements.append(Paragraph("Network Site Survey Report", title_style))
        elements.append(Spacer(1, 0.2*inch))
        
        # Survey Info Table
        survey_data = [
            ["Survey Name:", survey.get('name', 'N/A')],
            ["Client:", survey.get('client', 'N/A')],
            ["Site:", survey.get('site', 'N/A')],
            ["Description:", survey.get('description', 'N/A')],
            ["Created:", survey.get('created_at', 'N/A')]
        ]
        
        survey_table = Table(survey_data, colWidths=[1.5*inch, 4*inch])
        survey_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (0, -1), colors.HexColor('#f3f4f6')),
            ('TEXTCOLOR', (0, 0), (0, -1), colors.HexColor('#374151')),
            ('ALIGN', (0, 0), (0, -1), 'LEFT'),
            ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, -1), 10),
            ('GRID', (0, 0), (-1, -1), 1, colors.HexColor('#e5e7eb')),
            ('VALIGN', (0, 0), (-1, -1), 'TOP'),
            ('LEFTPADDING', (0, 0), (-1, -1), 8),
            ('RIGHTPADDING', (0, 0), (-1, -1), 8),
            ('TOPPADDING', (0, 0), (-1, -1), 6),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 6),
        ]))
        elements.append(survey_table)
        elements.append(Spacer(1, 0.3*inch))
        
        # Survey Responses Section
        elements.append(Paragraph("Survey Responses", heading_style))
        elements.append(Spacer(1, 0.1*inch))
        
        for q in SURVEY_TEMPLATE["questions"]:
            qid = q["id"]
            answer = survey.get("answers", {}).get(qid, "Not answered")
            
            # Question
            elements.append(Paragraph(f"<b>{q['question']}</b>", normal_style))
            # Answer
            elements.append(Paragraph(f"{answer}", normal_style))
            elements.append(Spacer(1, 0.1*inch))
        
        # AI Analysis Section
        if content:
            elements.append(PageBreak())
            elements.append(Paragraph("AI Analysis & Recommendations", heading_style))
            elements.append(Spacer(1, 0.1*inch))
            elements.append(Paragraph(content, normal_style))
        
        # Build PDF
        doc.build(elements)
        
        # Get PDF content
        pdf_content = buffer.getvalue()
        buffer.close()
        
        return jsonify({
            "format": "pdf",
            "filename": f"{survey.get('name', 'survey').replace(' ', '_')}_report.pdf",
            "content": base64.b64encode(pdf_content).decode('utf-8'),
            "mimeType": "application/pdf"
        })
    
    elif format_type == "email":
        # Generate email with full report inline (mailto can't attach files)
        subject = f"Network Site Survey Report - {survey.get('name', 'N/A')}"
        
        # Build email body with full survey data
        body = f"""Network Site Survey Report
{'=' * 50}

Survey Name: {survey.get('name', 'N/A')}
Client: {survey.get('client', 'N/A')}
Site: {survey.get('site', 'N/A')}
Description: {survey.get('description', 'N/A')}
Created: {survey.get('created_at', 'N/A')}

{'=' * 50}
SURVEY RESPONSES
{'=' * 50}
"""
        
        for q in SURVEY_TEMPLATE["questions"]:
            qid = q["id"]
            answer = survey.get("answers", {}).get(qid, "Not answered")
            body += f"\n\nQ: {q['question']}\nA: {answer}"
        
        if content:
            body += f"\n\n{'=' * 50}\nAI ANALYSIS & RECOMMENDATIONS\n{'=' * 50}\n\n{content}"
        
        body += "\n\nGenerated by BeawIT Network Site Survey AI"
        
        # Create mailto link (truncated if too long for URL)
        max_body_length = 1800  # Mailto links have length limits
        if len(body) > max_body_length:
            truncated_body = body[:max_body_length] + "\n\n[Report truncated due to email length limits. Use PDF or Text export for full report.]"
        else:
            truncated_body = body
        
        mailto_link = f"mailto:?subject={urllib.parse.quote(subject)}&body={urllib.parse.quote(truncated_body)}"
        
        return jsonify({
            "format": "email",
            "mailto_link": mailto_link,
            "subject": subject,
            "body": body
        })
    
    else:
        return jsonify({"error": "Unsupported format"}), 400

if __name__ == '__main__':
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 5000
    app.run(host='0.0.0.0', port=port, debug=True)
