src/app.py

"""
Flask application entry point.
Feature: 012-profile-contact-ui
"""

import json
import logging
import os
import time
from datetime import timedelta

from flask import Flask, jsonify, request, g
from flask_session import Session
from dotenv import load_dotenv
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator

from .services.auth_service import auth_service
from .services.storage_service import init_db
from .utils.tracing import init_tracer

# Load environment variables
load_dotenv()


class JSONFormatter(logging.Formatter):
    """Custom JSON formatter for structured logging."""

    def format(self, record):
        """Format log record as JSON."""
        log_obj = {
            "timestamp": self.formatTime(record, self.datefmt),
            "level": record.levelname,
            "module": record.module,
            "message": record.getMessage(),
        }

        # Add exception info if present
        if record.exc_info:
            log_obj["exception"] = self.formatException(record.exc_info)

        # Add extra fields if present
        if hasattr(record, "request_id"):
            log_obj["request_id"] = record.request_id
        if hasattr(record, "user_id"):
            log_obj["user_id"] = record.user_id
        if hasattr(record, "duration_ms"):
            log_obj["duration_ms"] = record.duration_ms
        if hasattr(record, "backend_latency_ms"):
            log_obj["backend_latency_ms"] = record.backend_latency_ms
        if hasattr(record, "status_code"):
            log_obj["status_code"] = record.status_code
        if hasattr(record, "method"):
            log_obj["method"] = record.method
        if hasattr(record, "path"):
            log_obj["path"] = record.path

        return json.dumps(log_obj)


def create_app():
    """Create and configure Flask application."""
    app = Flask(__name__)

    # Configuration
    app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "dev-secret-key-change-me")
    
    # CRITICAL: Use server-side sessions to avoid cookie size limits (4KB)
    # OAuth state tokens make cookies exceed size limit, causing state to be lost
    app.config["SESSION_TYPE"] = "filesystem"  # Store sessions server-side
    app.config["SESSION_FILE_DIR"] = "/app/data/flask_sessions"  # Session storage directory
    app.config["SESSION_PERMANENT"] = False  # Default to non-permanent (override per-session)
    app.config["SESSION_USE_SIGNER"] = True  # Sign session cookie for security
    
    # Session cookie configuration
    # For HTTPS (HF Spaces): Set SESSION_COOKIE_SECURE=true
    secure_cookie = os.getenv("SESSION_COOKIE_SECURE", "False").lower() == "true"
    app.config["SESSION_COOKIE_SECURE"] = secure_cookie
    app.config["SESSION_COOKIE_HTTPONLY"] = (
        os.getenv("SESSION_COOKIE_HTTPONLY", "True").lower() == "true"
    )
    # Use "None" for cross-site OAuth (HF Spaces), "Lax" for same-site (local dev)
    samesite_value = os.getenv("SESSION_COOKIE_SAMESITE", "Lax")
    app.config["SESSION_COOKIE_SAMESITE"] = None if samesite_value == "None" else samesite_value
    app.config["SESSION_COOKIE_NAME"] = "prepmate_session"  # Explicit session cookie name
    # Don't set SESSION_COOKIE_DOMAIN - let Flask handle it automatically
    app.config["SESSION_COOKIE_PATH"] = "/"  # Ensure cookie is valid for all paths
    app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(
        seconds=int(os.getenv("PERMANENT_SESSION_LIFETIME", "2592000"))
    )
    
    # CRITICAL: Log ALL environment variables for OAuth debugging
    print("="*80)
    print("[CONFIG] ENVIRONMENT VARIABLES CHECK")
    print(f"[CONFIG] SESSION_COOKIE_SECURE env var: {os.getenv('SESSION_COOKIE_SECURE', 'NOT SET')}")
    print(f"[CONFIG] SESSION_COOKIE_SAMESITE env var: {os.getenv('SESSION_COOKIE_SAMESITE', 'NOT SET')}")
    print(f"[CONFIG] PREFERRED_URL_SCHEME env var: {os.getenv('PREFERRED_URL_SCHEME', 'NOT SET')}")
    print(f"[CONFIG] SECRET_KEY env var: {'SET' if os.getenv('SECRET_KEY') else 'NOT SET'}")
    print("-"*80)
    print("[CONFIG] FLASK CONFIG VALUES:")
    print(f"[CONFIG] SESSION_COOKIE_SECURE={secure_cookie}")
    print(f"[CONFIG] SESSION_COOKIE_SAMESITE={app.config['SESSION_COOKIE_SAMESITE']}")
    print(f"[CONFIG] SESSION_COOKIE_NAME={app.config['SESSION_COOKIE_NAME']}")
    print(f"[CONFIG] PERMANENT_SESSION_LIFETIME={app.config['PERMANENT_SESSION_LIFETIME']}")
    print("="*80)
    
    # FAIL LOUDLY if OAuth won't work
    if not secure_cookie or app.config['SESSION_COOKIE_SAMESITE'] != None:
        print("⚠️  WARNING: OAuth will FAIL - cookie won't persist!")
        print("⚠️  Set these in HuggingFace Space Settings → Variables:")
        print("    SESSION_COOKIE_SECURE=true")
        print("    SESSION_COOKIE_SAMESITE=None")
    
    # Authlib configuration for OAuth state management
    app.config["AUTHLIB_INSECURE_TRANSPORT"] = os.getenv("FLASK_ENV") == "development"
    
    # Trust proxy headers for HTTPS detection (needed for HuggingFace Spaces)
    app.config["PREFERRED_URL_SCHEME"] = os.getenv("PREFERRED_URL_SCHEME", "http")
    
    # Configure ProxyFix middleware to trust X-Forwarded-* headers
    from werkzeug.middleware.proxy_fix import ProxyFix
    app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1)

    # Initialize database
    init_db()

    # Initialize Flask-Session (must be BEFORE OAuth init)
    # This stores sessions server-side to avoid 4KB cookie size limit
    print("="*80)
    print("[FLASK-SESSION] INITIALIZING SERVER-SIDE SESSIONS")
    print(f"[FLASK-SESSION] SESSION_TYPE: {app.config.get('SESSION_TYPE')}")
    print(f"[FLASK-SESSION] SESSION_FILE_DIR: {app.config.get('SESSION_FILE_DIR')}")
    print(f"[FLASK-SESSION] SESSION_PERMANENT: {app.config.get('SESSION_PERMANENT')}")
    print(f"[FLASK-SESSION] SESSION_USE_SIGNER: {app.config.get('SESSION_USE_SIGNER')}")
    
    # Check if session directory exists
    session_dir = app.config.get('SESSION_FILE_DIR')
    if session_dir and os.path.exists(session_dir):
        print(f"[FLASK-SESSION] Session directory exists: {session_dir}")
        print(f"[FLASK-SESSION] Directory writable: {os.access(session_dir, os.W_OK)}")
    else:
        print(f"[FLASK-SESSION] ⚠️ WARNING: Session directory does NOT exist: {session_dir}")
    
    Session(app)
    
    # Verify Flask-Session is active
    print(f"[FLASK-SESSION] Session interface type: {type(app.session_interface).__name__}")
    print(f"[FLASK-SESSION] Session interface module: {type(app.session_interface).__module__}")
    if 'flask_session' in type(app.session_interface).__module__:
        print("[FLASK-SESSION] ✅ Server-side sessions ACTIVE")
    else:
        print("[FLASK-SESSION] ❌ ERROR: Still using client-side sessions!")
    print("="*80)

    # Initialize OAuth
    auth_service.init_app(app)

    # Initialize Jaeger tracing
    init_tracer("prepmate-webapp")

    # Configure logging
    setup_logging(app)

    # Register middleware
    register_middleware(app)

    # Register blueprints
    from .routes import auth, profile, contacts, settings

    app.register_blueprint(auth.bp)
    app.register_blueprint(profile.bp)
    app.register_blueprint(contacts.contacts_bp)
    app.register_blueprint(settings.bp)

    # Root route - redirect to login
    @app.route("/")
    def index():
        """Root route - redirect to login page."""
        from flask import session, redirect, url_for, render_template
        
        if "user_id" in session:
            return redirect(url_for("profile.view_profile"))
        return render_template("login.html")

    # Health check endpoint
    @app.route("/health")
    def health():
        """Health check endpoint for monitoring."""
        return jsonify({"status": "ok"}), 200
    
    # Diagnostic endpoint to check environment variables
    @app.route("/debug/config")
    def debug_config():
        """Show current configuration for debugging."""
        return jsonify({
            "environment_variables": {
                "SESSION_COOKIE_SECURE": os.getenv("SESSION_COOKIE_SECURE", "NOT SET"),
                "SESSION_COOKIE_SAMESITE": os.getenv("SESSION_COOKIE_SAMESITE", "NOT SET"),
                "PREFERRED_URL_SCHEME": os.getenv("PREFERRED_URL_SCHEME", "NOT SET"),
                "SECRET_KEY": "SET" if os.getenv("SECRET_KEY") else "NOT SET",
            },
            "flask_config": {
                "SESSION_COOKIE_SECURE": app.config.get("SESSION_COOKIE_SECURE"),
                "SESSION_COOKIE_SAMESITE": str(app.config.get("SESSION_COOKIE_SAMESITE")),
                "SESSION_COOKIE_NAME": app.config.get("SESSION_COOKIE_NAME"),
                "SESSION_COOKIE_PATH": app.config.get("SESSION_COOKIE_PATH"),
                "PREFERRED_URL_SCHEME": app.config.get("PREFERRED_URL_SCHEME"),
            }
        }), 200

    return app


def setup_logging(app):
    """Configure structured JSON logging."""
    log_level = os.getenv("LOG_LEVEL", "INFO")

    # Create handler with JSON formatter
    handler = logging.StreamHandler()
    handler.setFormatter(JSONFormatter())

    # Configure root logger
    logging.basicConfig(
        level=getattr(logging, log_level),
        handlers=[handler],
    )

    # Configure app logger
    app.logger.handlers = [handler]
    app.logger.setLevel(getattr(logging, log_level))


def register_middleware(app):
    """Register Flask middleware for request logging and timing."""

    @app.before_request
    def before_request():
        """Start request timer and generate request ID, create tracing span."""
        from .utils.tracing import is_tracing_enabled
        
        g.start_time = time.time()
        g.request_id = request.headers.get("X-Request-ID", os.urandom(8).hex())
        
        # Skip tracing if disabled or for health endpoint
        if not is_tracing_enabled() or request.path == "/health":
            return
        
        # Extract parent span context from headers if present
        try:
            ctx = TraceContextTextMapPropagator().extract(carrier=request.headers)
        except Exception:
            ctx = None
        
        # Start a new span for this request
        tracer = trace.get_tracer(__name__)
        span = tracer.start_span(
            name=f"{request.method} {request.path}",
            context=ctx,
        )
        span.set_attribute("http.method", request.method)
        span.set_attribute("http.url", request.url)
        span.set_attribute("request_id", g.request_id)
        g.span = span

    @app.after_request
    def after_request(response):
        """Log request completion with duration and finish span."""
        if hasattr(g, "start_time"):
            duration_ms = (time.time() - g.start_time) * 1000

            # Skip logging for Streamlit internal requests unless in debug mode
            is_streamlit_internal = request.path.startswith("/_stcore")
            is_debug_mode = app.logger.level == logging.DEBUG
            
            if not is_streamlit_internal or is_debug_mode:
                # Get user_id from session if available
                user_id = None
                try:
                    from flask import session

                    user_id = session.get("user_id")
                except Exception:
                    pass

                # Log request with structured data
                extra = {
                    "request_id": g.request_id,
                    "duration_ms": round(duration_ms, 2),
                    "status_code": response.status_code,
                    "method": request.method,
                    "path": request.path,
                }
                if user_id:
                    extra["user_id"] = user_id
                if hasattr(g, "backend_latency_ms"):
                    extra["backend_latency_ms"] = round(g.backend_latency_ms, 2)

                app.logger.info(
                    f"{request.method} {request.path} {response.status_code}",
                    extra=extra,
                )
        
        # Finish tracing span
        if hasattr(g, "span"):
            span = g.span
            span.set_attribute("http.status_code", response.status_code)
            if response.status_code >= 400:
                span.set_status(Status(StatusCode.ERROR))
            if hasattr(g, "backend_latency_ms"):
                span.set_attribute("backend_latency_ms", round(g.backend_latency_ms, 2))
            span.end()

        return response
    
    # Session debugging and cookie SameSite=None fix
    @app.after_request
    def debug_and_fix_session_cookie(response):
        """Debug session cookie setting and force SameSite=None if needed."""
        from flask import session
        
        # CRITICAL FIX: Force Flask-Session to set cookie when session is modified
        # Flask-Session only sets cookies for new sessions or when session.permanent changes
        # But we need it to ALWAYS set a cookie when session is modified to replace old client-side cookies
        if session.modified and not response.headers.get('Set-Cookie'):
            # Manually call save_session to force cookie setting
            app.session_interface.save_session(app, session, response)
        
        # ALWAYS log, even if session is empty
        print(f"[AFTER-REQUEST] Path: {request.path}, Status: {response.status_code}")
        print(f"[AFTER-REQUEST] Session modified: {session.modified}, Has data: {bool(session)}, Keys: {list(session.keys())}")
        print(f"[AFTER-REQUEST] Session interface: {type(app.session_interface).__name__}")
        
        # Log session state and response cookies
        set_cookie_headers = response.headers.getlist('Set-Cookie')
        has_session_cookie = any('prepmate_session' in cookie for cookie in set_cookie_headers)
        
        print(f"[AFTER-REQUEST] Set-Cookie count: {len(set_cookie_headers)}")
        print(f"[AFTER-REQUEST] Has prepmate_session cookie: {has_session_cookie}")
        print(f"[AFTER-REQUEST] All Set-Cookie headers: {set_cookie_headers}")
        
        if has_session_cookie:
            # Check if cookie looks like session ID (server-side) or full session (client-side)
            for cookie in set_cookie_headers:
                if 'prepmate_session' in cookie:
                    # Extract cookie value (between = and ;)
                    cookie_value = cookie.split('=')[1].split(';')[0] if '=' in cookie else ''
                    print(f"[AFTER-REQUEST] Cookie value length: {len(cookie_value)} chars")
                    print(f"[AFTER-REQUEST] Cookie value preview: {cookie_value[:50]}...")
                    if cookie_value.startswith('.'):
                        print("[AFTER-REQUEST] ⚠️ Cookie looks like CLIENT-SIDE session (starts with '.')")
                    elif len(cookie_value) < 100:
                        print("[AFTER-REQUEST] ✅ Cookie looks like SERVER-SIDE session ID (short)")
                    else:
                        print("[AFTER-REQUEST] ⚠️ Cookie looks like CLIENT-SIDE session (long)")
            
            # CRITICAL FIX: Flask's session interface doesn't properly set SameSite=None and Partitioned
            # Manually fix any prepmate_session cookies that are missing these attributes
            print("[AFTER-REQUEST] Fixing SameSite=None and Partitioned for session cookie")
            new_cookies = []
            for cookie in set_cookie_headers:
                if 'prepmate_session' in cookie:
                    # Check if SameSite=None is missing
                    if 'SameSite=None' not in cookie:
                        cookie = cookie.rstrip(';') + '; SameSite=None'
                        print(f"[AFTER-REQUEST] Added SameSite=None to session cookie")
                    
                    # Check if Partitioned is missing (required for CHIPS - Cookies Having Independent Partitioned State)
                    if 'Partitioned' not in cookie:
                        cookie = cookie.rstrip(';') + '; Partitioned'
                        print(f"[AFTER-REQUEST] Added Partitioned to session cookie")
                    
                    new_cookies.append(cookie)
                else:
                    # Keep non-session cookies unchanged
                    new_cookies.append(cookie)
            
            # Replace all Set-Cookie headers
            response.headers.remove('Set-Cookie')
            for cookie in new_cookies:
                response.headers.add('Set-Cookie', cookie)
            
            print(f"[AFTER-REQUEST] Final Set-Cookie headers: {response.headers.getlist('Set-Cookie')}")
        else:
            print("[AFTER-REQUEST] ⚠️ WARNING: No session cookie found in response!")
            print(f"[AFTER-REQUEST] This means Flask-Session did not set a cookie")
            print(f"[AFTER-REQUEST] Session was modified: {session.modified}")
            print(f"[AFTER-REQUEST] Session is new: {getattr(session, 'new', 'N/A')}")
        
        return response


# Create app instance
app = create_app()

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.getenv("PORT", "5000")), debug=True)