#!/usr/bin/env python3
"""
IT Site Survey AI Application
Complete implementation with all requested features:
- Photo uploads for AI assessment
- Model selection (Ollama vision models)
- Export results (PDF, Text, Email)
- Delete surveys (individual + Clear All)
- Chat interface for AI interaction
- Load Sample Data for testing
- Comprehensive 35-question survey template
"""

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import os
import json
import uuid
from datetime import datetime
import base64
import requests
from werkzeug.utils import secure_filename

app = Flask(__name__)
CORS(app, resources={
    r"/api/*": {
        "origins": ["http://192.168.50.11:3003", "http://localhost:3003"],
        "methods": ["GET", "POST", "PUT", "DELETE", "OPTIONS"],
        "allow_headers": ["Content-Type", "Authorization"]
    }
})

# Configuration
UPLOAD_FOLDER = 'uploads'
PHOTOS_FOLDER = 'uploads/photos'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://192.168.19.25:11434')
OLLAMA_URL = f"{OLLAMA_HOST}/api/chat"

# Create upload directories
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(PHOTOS_FOLDER, exist_ok=True)

# In-memory storage
surveys_db = {}
responses_db = {}

# Survey Template with 35 comprehensive questions
DEFAULT_SURVEY_TEMPLATE = {
    "title": "IT Infrastructure Site Survey",
    "description": "Comprehensive network infrastructure assessment",
    "questions": [
        # Site Topology & Layout (1-4)
        {
            "id": "q1",
            "category": "Site Topology & Layout",
            "question": "What is the total square footage of the site?",
            "type": "number",
            "placeholder": "e.g., 3500"
        },
        {
            "id": "q2",
            "category": "Site Topology & Layout",
            "question": "How many floors or buildings need network coverage?",
            "type": "number",
            "placeholder": "e.g., 2"
        },
        {
            "id": "q3",
            "category": "Site Topology & Layout",
            "question": "Describe the site layout (open plan, offices, warehouse, etc.)",
            "type": "textarea",
            "placeholder": "Describe the layout..."
        },
        {
            "id": "q4",
            "category": "Site Topology & Layout",
            "question": "Are there any physical barriers that may affect WiFi coverage?",
            "type": "multiselect",
            "options": ["Concrete walls", "Metal shelving", "Elevator shafts", "Thick drywall", "None"]
        },
        
        # Network Closet & Equipment (5-8)
        {
            "id": "q5",
            "category": "Network Closet & Equipment",
            "question": "How many network closets or IDF rooms are on-site?",
            "type": "number",
            "placeholder": "e.g., 2"
        },
        {
            "id": "q6",
            "category": "Network Closet & Equipment",
            "question": "Describe the current network closet setup",
            "type": "textarea",
            "placeholder": "Describe the setup..."
        },
        {
            "id": "q7",
            "category": "Network Closet & Equipment",
            "question": "What type of power and cooling is available in the network closets?",
            "type": "textarea",
            "placeholder": "Describe power and cooling..."
        },
        {
            "id": "q8",
            "category": "Network Closet & Equipment",
            "question": "Are there existing server racks or wall-mounted equipment?",
            "type": "select",
            "options": ["Yes, full racks", "Yes, partial racks", "Wall-mounted only", "No existing racks"]
        },
        
        # Cabling Infrastructure (9-12)
        {
            "id": "q9",
            "category": "Cabling Infrastructure",
            "question": "What type of existing cabling is in place?",
            "type": "select",
            "options": ["Cat5e (100Mb-1GbE)", "Cat6 (1GbE-10GbE)", "Cat6a (10GbE)", "Fiber (Multi-mode)", "Fiber (Single-mode)", "No existing cabling", "Unknown"]
        },
        {
            "id": "q10",
            "category": "Cabling Infrastructure",
            "question": "How many network drops are currently installed?",
            "type": "number",
            "placeholder": "e.g., 120"
        },
        {
            "id": "q11",
            "category": "Cabling Infrastructure",
            "question": "What is the estimated distance from the MDF to the farthest IDF or workstation?",
            "type": "textarea",
            "placeholder": "Describe distances..."
        },
        {
            "id": "q12",
            "category": "Cabling Infrastructure",
            "question": "Are there any special cabling requirements?",
            "type": "multiselect",
            "options": ["Plenum-rated cable required", "Outdoor-rated cable needed", "Direct burial cable needed", "Aerial cable needed", "Shielded cable required", "None"]
        },
        
        # Network Equipment (13-15)
        {
            "id": "q13",
            "category": "Network Equipment",
            "question": "What network equipment is currently installed?",
            "type": "multiselect",
            "options": ["Router/Firewall", "Managed switches", "Unmanaged switches", "Wireless access points", "None/Unknown"]
        },
        {
            "id": "q14",
            "category": "Network Equipment",
            "question": "What type of switches are needed?",
            "type": "select",
            "options": ["Unmanaged (plug-and-play)", "Managed (Layer 2)", "Managed (Layer 3 with routing)", "PoE switches", "Industrial-grade switches"]
        },
        {
            "id": "q15",
            "category": "Network Equipment",
            "question": "How many wireless access points are needed?",
            "type": "number",
            "placeholder": "e.g., 6"
        },
        
        # Power Requirements (16-18)
        {
            "id": "q16",
            "category": "Power Requirements",
            "question": "What is the estimated total power consumption for all network equipment?",
            "type": "textarea",
            "placeholder": "Describe power requirements..."
        },
        {
            "id": "q17",
            "category": "Power Requirements",
            "question": "What PoE devices will be connected?",
            "type": "multiselect",
            "options": ["IP Phones", "Wireless APs", "IP Cameras", "IoT sensors", "None", "Not sure"]
        },
        {
            "id": "q18",
            "category": "Power Requirements",
            "question": "What UPS backup capacity is required?",
            "type": "select",
            "options": ["No backup needed", "Minimal backup (15-30 min)", "Standard backup (1-2 hours)", "Extended backup (4+ hours)", "Generator backup"]
        },
        
        # User & Device Requirements (19-21)
        {
            "id": "q19",
            "category": "User & Device Requirements",
            "question": "How many employees or users will use the network?",
            "type": "number",
            "placeholder": "e.g., 45"
        },
        {
            "id": "q20",
            "category": "User & Device Requirements",
            "question": "What is the estimated total number of devices?",
            "type": "number",
            "placeholder": "e.g., 85"
        },
        {
            "id": "q21",
            "category": "User & Device Requirements",
            "question": "What types of network activities will be performed?",
            "type": "multiselect",
            "options": ["Email/Web Browsing", "Video conferencing", "File sharing", "Cloud applications", "Remote desktop", "VoIP/Phone calls", "Streaming media"]
        },
        
        # Security & Compliance (22-24)
        {
            "id": "q22",
            "category": "Security & Compliance",
            "question": "What compliance standards must be met?",
            "type": "multiselect",
            "options": ["HIPAA", "PCI-DSS", "SOC 2", "GDPR", "NIST", "None/General security"]
        },
        {
            "id": "q23",
            "category": "Security & Compliance",
            "question": "What network segmentation is required?",
            "type": "multiselect",
            "options": ["Corporate network", "Guest network", "IoT network", "Management network", "No segmentation needed"]
        },
        {
            "id": "q24",
            "category": "Security & Compliance",
            "question": "What security features are needed?",
            "type": "multiselect",
            "options": ["Firewall", "VPN access", "Intrusion detection", "Network access control (NAC)", "Content filtering", "None"]
        },
        
        # Internet & WAN (25-27)
        {
            "id": "q25",
            "category": "Internet & WAN",
            "question": "What type of internet connectivity is available?",
            "type": "multiselect",
            "options": ["Fiber", "Cable/DSL", "5G/Cellular backup", "Dedicated leased line", "No existing connection"]
        },
        {
            "id": "q26",
            "category": "Internet & WAN",
            "question": "What is the required internet bandwidth?",
            "type": "textarea",
            "placeholder": "e.g., 500Mbps download / 250Mbps upload"
        },
        {
            "id": "q27",
            "category": "Internet & WAN",
            "question": "Are there multiple sites that need to be connected?",
            "type": "select",
            "options": ["No - Single site only", "Yes - 2-3 sites", "Yes - 4+ sites", "Not sure yet"]
        },
        
        # Existing Infrastructure (28-30)
        {
            "id": "q28",
            "category": "Existing Infrastructure",
            "question": "Is there existing IT infrastructure to integrate with?",
            "type": "select",
            "options": ["Yes - Full upgrade needed", "Yes - Partial upgrade", "No - New installation", "Not sure"]
        },
        {
            "id": "q29",
            "category": "Existing Infrastructure",
            "question": "Describe any existing network issues or limitations",
            "type": "textarea",
            "placeholder": "Describe current issues..."
        },
        {
            "id": "q30",
            "category": "Existing Infrastructure",
            "question": "What performance issues have been observed?",
            "type": "multiselect",
            "options": ["Slow internet", "WiFi dead zones", "Network congestion", "Frequent disconnections", "No issues observed"]
        },
        
        # Budget & Timeline (31-33)
        {
            "id": "q31",
            "category": "Budget & Timeline",
            "question": "What is the estimated budget for this project?",
            "type": "select",
            "options": ["Under $10,000", "$10,000 - $25,000", "$25,000 - $50,000", "$50,000 - $100,000", "$100,000+", "Not determined yet"]
        },
        {
            "id": "q32",
            "category": "Budget & Timeline",
            "question": "What is the desired timeline for completion?",
            "type": "select",
            "options": ["Within 1 month", "Within 3 months", "Within 6 months", "Within 1 year", "Flexible"]
        },
        {
            "id": "q33",
            "category": "Budget & Timeline",
            "question": "Who will be responsible for ongoing network management?",
            "type": "select",
            "options": ["Internal IT team", "Managed service provider (MSP)", "Hybrid approach", "Not determined yet"]
        },
        
        # Additional Details (34-35)
        {
            "id": "q34",
            "category": "Additional Details",
            "question": "Upload any site photos, floor plans, or network diagrams",
            "type": "file",
            "placeholder": "Upload photos..."
        },
        {
            "id": "q35",
            "category": "Additional Details",
            "question": "Any additional requirements or special considerations?",
            "type": "textarea",
            "placeholder": "Describe any additional requirements..."
        }
    ]
}

# Sample data for testing
SAMPLE_SURVEY_DATA = {
    "name": "Sample Corporate Office Survey",
    "client_name": "Acme Corporation",
    "site_name": "Main Office - Downtown",
    "description": "Comprehensive network infrastructure assessment for downtown office location",
    "responses": {
        "q1": "3500",
        "q2": "2",
        "q3": "Single-story office building with open floor plan, 8 private offices, 2 conference rooms, break room, and reception area. Approximately 60 workstations in cubicle area.",
        "q4": ["Concrete walls"],
        "q5": "2",
        "q6": "MDF located in basement utility room. IDF on first floor in IT storage closet. Both have adequate ventilation but limited space.",
        "q7": "2x 20A circuits per closet, UPS backup planned for MDF",
        "q8": "Yes, partial racks",
        "q9": "Cat5e (100Mb-1GbE)",
        "q10": "120",
        "q11": "Average 150 feet, longest run approximately 280 feet to corner offices",
        "q12": ["Plenum-rated cable required"],
        "q13": ["Router/Firewall", "Unmanaged switches", "Wireless access points"],
        "q14": "Managed (Layer 2)",
        "q15": "6",
        "q16": "1800W estimated for all network equipment including PoE devices",
        "q17": ["IP Phones", "Wireless APs", "IP Cameras"],
        "q18": "Limited backup (1-2 hours)",
        "q19": "45",
        "q20": "85",
        "q21": ["Email/Web Browsing", "Video conferencing", "File sharing", "Cloud applications", "Remote desktop"],
        "q22": ["SOC 2"],
        "q23": ["Corporate network", "Guest network", "IoT network", "Management network"],
        "q24": ["Firewall", "VPN access", "Network access control (NAC)"],
        "q25": ["Fiber", "5G/Cellular backup"],
        "q26": "500Mbps download / 250Mbps upload",
        "q27": "No - Single site only",
        "q28": "Yes - Partial upgrade",
        "q29": "Current WiFi has dead zones in conference rooms. Existing switches are unmanaged and showing signs of age. Need better coverage and management capabilities.",
        "q30": ["Slow internet", "WiFi dead zones", "Network congestion"],
        "q31": "$25,000 - $50,000",
        "q32": "Within 3 months",
        "q33": "Internal IT team",
        "q34": "",
        "q35": "Need outdoor WiFi coverage for patio area. Plan to add IP cameras in future. Prefer Cisco or Ubiquiti equipment. Require weekend installation to minimize downtime."
    }
}


@app.route('/api/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    try:
        resp = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        ollama_online = resp.status_code == 200
    except:
        ollama_online = False
    
    return jsonify({
        "status": "ok",
        "service": "site-survey-ai",
        "timestamp": datetime.utcnow().isoformat(),
        "ollama_connected": ollama_online
    })


@app.route('/api/models', methods=['GET'])
def get_models():
    """Get available Ollama models"""
    try:
        resp = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        if resp.status_code == 200:
            models_data = resp.json()
            models = [m['name'] for m in models_data.get('models', [])]
            # Add cloud models if available
            cloud_models = ['kimi-k2.6:cloud', 'kimi-k2.5:cloud', 'nemotron-3-super:cloud']
            models.extend(cloud_models)
            return jsonify({"models": models, "default": "kimi-k2.6:cloud"})
    except Exception as e:
        print(f"Error fetching models: {e}")
    
    return jsonify({
        "models": ["kimi-k2.6:cloud", "kimi-k2.5:cloud", "llava:latest"],
        "default": "kimi-k2.6:cloud"
    })


@app.route('/api/ollama-status', methods=['GET'])
def ollama_status():
    """Check Ollama connection status"""
    try:
        resp = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        return jsonify({"online": resp.status_code == 200})
    except:
        return jsonify({"online": False})


@app.route('/api/surveys/template', methods=['GET'])
def get_survey_template():
    """Get the survey template with questions"""
    return jsonify({"template": DEFAULT_SURVEY_TEMPLATE})


@app.route('/api/surveys', methods=['GET'])
def list_surveys():
    """List all surveys"""
    surveys = []
    for sid, survey in surveys_db.items():
        surveys.append({
            "id": sid,
            "name": survey['name'],
            "client_name": survey.get('client_name', ''),
            "site_name": survey.get('site_name', ''),
            "description": survey.get('description', ''),
            "created_at": survey['created_at'],
            "responses_count": len(survey.get('responses', []))
        })
    return jsonify({"surveys": surveys})


@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', 'New Survey'),
        'description': data.get('description', ''),
        'client_name': data.get('client_name', ''),
        'site_name': data.get('site_name', ''),
        'questions': DEFAULT_SURVEY_TEMPLATE['questions'],
        'created_at': datetime.utcnow().isoformat(),
        'status': 'active',
        'responses': []
    }
    
    # If responses are provided, add them
    responses_data = data.get('responses')
    if responses_data and isinstance(responses_data, dict) and len(responses_data) > 0:
        response_entry = {
            'id': str(uuid.uuid4()),
            'survey_id': survey_id,
            'responses': responses_data,
            'photos': [],
            'analyzed': False,
            'analysis': None,
            'submitter_name': 'Survey Creator',
            'created_at': datetime.utcnow().isoformat()
        }
        survey['responses'].append(response_entry)
    
    surveys_db[survey_id] = survey
    return jsonify({"survey": survey}), 201


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


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


@app.route('/api/surveys', methods=['DELETE'])
def delete_all_surveys():
    """Delete all surveys"""
    count = len(surveys_db)
    surveys_db.clear()
    return jsonify({"success": True, "message": f"Deleted {count} surveys"})


@app.route('/api/surveys/<survey_id>/responses', methods=['POST'])
def submit_response(survey_id):
    """Submit survey responses"""
    if survey_id not in surveys_db:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json()
    response_id = str(uuid.uuid4())
    
    response_entry = {
        'id': response_id,
        'survey_id': survey_id,
        'responses': data.get('responses', {}),
        'photos': data.get('photos', []),
        'analyzed': False,
        'analysis': None,
        'submitter_name': data.get('submitter_name', 'Anonymous'),
        'created_at': datetime.utcnow().isoformat()
    }
    
    surveys_db[survey_id]['responses'].append(response_entry)
    return jsonify({"response": response_entry}), 201


@app.route('/api/surveys/<survey_id>/photos', methods=['POST'])
def upload_photo(survey_id):
    """Upload a photo for a survey"""
    if survey_id not in surveys_db:
        return jsonify({"error": "Survey not found"}), 404
    
    if 'photo' not in request.files:
        return jsonify({"error": "No photo provided"}), 400
    
    photo = request.files['photo']
    if photo.filename == '':
        return jsonify({"error": "No photo selected"}), 400
    
    # Save photo
    filename = secure_filename(f"{survey_id}_{uuid.uuid4()}_{photo.filename}")
    photo_path = os.path.join(PHOTOS_FOLDER, filename)
    photo.save(photo_path)
    
    return jsonify({
        "success": True,
        "filename": filename,
        "url": f"/api/photos/{survey_id}/{filename}"
    })


@app.route('/api/photos/<survey_id>/<filename>', methods=['GET'])
def get_photo(survey_id, filename):
    """Get a photo"""
    photo_path = os.path.join(PHOTOS_FOLDER, filename)
    if os.path.exists(photo_path):
        return send_file(photo_path)
    return jsonify({"error": "Photo not found"}), 404


@app.route('/api/surveys/<survey_id>/analyze', methods=['POST'])
def analyze_survey(survey_id):
    """Analyze survey with AI"""
    if survey_id not in surveys_db:
        return jsonify({"error": "Survey not found"}), 404
    
    survey = surveys_db[survey_id]
    data = request.get_json() or {}
    selected_model = data.get('model', 'llava:latest')
    
    # Get the latest response
    if not survey['responses']:
        return jsonify({"error": "No responses to analyze"}), 400
    
    last_response = survey['responses'][-1]
    responses = last_response['responses']
    
    # Build analysis prompt
    prompt = f"""Analyze this IT site survey and provide recommendations:

Survey: {survey['name']}
Client: {survey['client_name']}
Site: {survey['site_name']}

Responses:
"""
    
    for q in survey['questions']:
        qid = q['id']
        if qid in responses:
            prompt += f"\n{q['question']}: {responses[qid]}"
    
    prompt += "\n\nProvide a detailed analysis including:\n"
    prompt += "1. Network infrastructure recommendations\n"
    prompt += "2. Equipment recommendations with estimated costs\n"
    prompt += "3. Implementation timeline\n"
    prompt += "4. Potential challenges and solutions\n"
    
    try:
        # Call Ollama for analysis
        ollama_data = {
            "model": selected_model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": False
        }
        
        resp = requests.post(OLLAMA_URL, json=ollama_data, timeout=300)
        if resp.status_code == 200:
            result = resp.json()
            analysis = result.get('message', {}).get('content', 'No analysis available')
            
            # Save analysis
            last_response['analyzed'] = True
            last_response['analysis'] = analysis
            last_response['model_used'] = selected_model
            
            return jsonify({
                "analysis": analysis,
                "model_used": selected_model
            })
        else:
            return jsonify({"error": f"Ollama error: {resp.status_code}"}), 500
            
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/api/surveys/<survey_id>/export', methods=['POST'])
def export_survey(survey_id):
    """Export survey results"""
    if survey_id not in surveys_db:
        return jsonify({"error": "Survey not found"}), 404
    
    data = request.get_json() or {}
    export_format = data.get('format', 'json')
    
    survey = surveys_db[survey_id]
    
    if export_format == 'json':
        return jsonify({"survey": survey})
    elif export_format == 'text':
        # Generate text report
        text = f"IT Site Survey Report\n"
        text += f"Survey: {survey['name']}\n"
        text += f"Client: {survey['client_name']}\n"
        text += f"Site: {survey['site_name']}\n"
        text += f"Date: {survey['created_at']}\n\n"
        
        for response in survey['responses']:
            text += f"Response by: {response['submitter_name']}\n"
            for q in survey['questions']:
                qid = q['id']
                if qid in response['responses']:
                    text += f"{q['question']}: {response['responses'][qid]}\n"
            if response.get('analysis'):
                text += f"\nAI Analysis:\n{response['analysis']}\n"
        
        return jsonify({"text": text})
    else:
        return jsonify({"error": "Unsupported format"}), 400


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=3003, debug=False)