"""
Sales Voice Gateway API

Production-ready API with:
- Organization API key authentication
- Tiered rate limiting
- Security middleware
- Latency optimizations
- Usage tracking for billing
"""
from __future__ import annotations

import os
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from fastapi import Depends, FastAPI, Header, HTTPException, Request, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from openai import APIConnectionError, APIError
from pydantic import BaseModel, Field

try:
    from dotenv import load_dotenv
except Exception:
    load_dotenv = None

from app.auth.jwt import AuthError, issue_access_token, issue_voice_token, verify_token
from app.auth.api_keys import (
    ApiKeyInfo,
    Plan,
    PLAN_CONFIGS,
    generate_api_key,
    validate_api_key,
    revoke_api_key,
)
from app.auth.rate_limiter import (
    RateLimitError,
    ConcurrencyLimitError,
    check_rate_limit,
    reserve_session,
    release_session,
    record_voice_usage,
    get_usage_summary,
    get_remaining_quota,
)
from app.auth.security import SecurityMiddleware, IPAllowlistMiddleware, get_client_ip
from app.storage.redis import get_redis, close_redis
from app.storage.database import init_pool, close_pool
from app.voice.ws import voice_ws_endpoint
from app.observability.logging import configure_logging, get_logger
from app.sales_brain.brain import SalesBrainEngine, merge_memory
from app.utils.latency import (
    warm_up_connections,
    get_all_latency_stats,
    timed,
    get_latency_tracker,
)

if load_dotenv:
    load_dotenv()

configure_logging()
_logger = get_logger("main")


# =============================================================================
# Lifespan Management
# =============================================================================

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler for startup/shutdown."""
    _logger.info("starting_up")
    
    # Initialize database pool
    try:
        await init_pool()
        _logger.info("database_pool_initialized")
    except Exception as e:
        _logger.error("database_pool_init_failed", error=str(e))
        # Continue without database - graceful degradation
    
    # Warm up connections on startup
    await warm_up_connections(get_redis)
    
    yield
    
    # Cleanup on shutdown
    _logger.info("shutting_down")
    await close_redis()
    await close_pool()


# =============================================================================
# Application Setup
# =============================================================================

app = FastAPI(
    title="Sales Voice Gateway API",
    description="Production-ready voice AI API for sales conversations",
    version="2.0.0",
    lifespan=lifespan,
    docs_url="/docs" if os.getenv("ENABLE_DOCS", "true").lower() == "true" else None,
    redoc_url="/redoc" if os.getenv("ENABLE_DOCS", "true").lower() == "true" else None,
)

# Add middleware (order matters - first added is outermost)
app.add_middleware(GZipMiddleware, minimum_size=1000)

# CORS configuration
allowed_origins = os.getenv("CORS_ORIGINS", "*").split(",")
app.add_middleware(
    CORSMiddleware,
    allow_origins=allowed_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=[
        "X-Request-ID",
        "X-Response-Time",
        "X-RateLimit-Limit-Minute",
        "X-RateLimit-Remaining-Minute",
        "X-RateLimit-Limit-Hour",
        "X-RateLimit-Remaining-Hour",
        "Retry-After",
    ],
)

# Security middleware
app.add_middleware(SecurityMiddleware, redis_factory=get_redis)

# Optional IP allowlist
allowed_ips = os.getenv("ALLOWED_IPS")
if allowed_ips:
    app.add_middleware(IPAllowlistMiddleware, allowed_ips=allowed_ips.split(","))

# Static files
WEB_ROOT = os.path.join(os.path.dirname(__file__), "webui")
app.mount("/static", StaticFiles(directory=WEB_ROOT), name="static")


# =============================================================================
# Request/Response Models
# =============================================================================

class CreateApiKeyRequest(BaseModel):
    org_id: str = Field(..., description="Organization identifier")
    plan: str = Field(default="free", description="Pricing plan tier")
    allowed_ips: Optional[List[str]] = Field(default=None, description="IP allowlist")
    webhook_url: Optional[str] = Field(default=None, description="Webhook for events")


class CreateApiKeyResponse(BaseModel):
    api_key: str = Field(..., description="API key (show only once)")
    key_id: str = Field(..., description="Key identifier for management")
    org_id: str
    plan: str
    limits: Dict[str, Any]


class RevokeApiKeyRequest(BaseModel):
    key_id: str = Field(..., description="Key ID to revoke")


class VoiceSessionRequest(BaseModel):
    format: Optional[str] = "pcm16"
    sample_rate: Optional[int] = 16000


class VoiceSessionResponse(BaseModel):
    sid: str
    token: str
    ws_url: str
    expires_in: int


class TokenRequest(BaseModel):
    user_id: str
    org: Optional[str] = None
    plan: str = "free"


class ChatRequest(BaseModel):
    session_id: Optional[str] = None
    message: str


class ChatResponse(BaseModel):
    session_id: str
    reply: str
    state: str
    lead_score: int
    missing_fields: List[str]
    memory: Dict[str, Any]
    safety_flags: List[str]
    turns: int


class UsageResponse(BaseModel):
    org_id: str
    requests_minute: int
    requests_hour: int
    requests_day: int
    concurrent_sessions: int
    total_voice_seconds: float
    limits: Dict[str, Any]


class HealthResponse(BaseModel):
    status: str
    version: str
    uptime_seconds: float
    latency_stats: List[Dict[str, Any]]


# =============================================================================
# Plan Configuration
# =============================================================================

PLAN_CONCURRENCY = {
    "free": 1,
    "starter": 3,
    "pro": 10,
    "enterprise": 50,
}


# =============================================================================
# Session Management
# =============================================================================

@dataclass
class ChatSession:
    sales_state: str = "GREETING"
    memory: Dict[str, Any] = field(default_factory=dict)
    turns: int = 0


_chat_sessions: Dict[str, ChatSession] = {}
_chat_brain = SalesBrainEngine()
_start_time = None


# =============================================================================
# Helper Functions
# =============================================================================

def _ws_url_from_request(request: Request) -> str:
    configured = os.getenv("VOICE_WS_URL")
    if configured:
        return configured
    host = request.headers.get("host", "localhost")
    scheme = "wss" if request.url.scheme == "https" else "ws"
    return f"{scheme}://{host}/voice"


def _strip_meta(memory: Dict[str, Any]) -> Dict[str, Any]:
    return {k: v for k, v in (memory or {}).items() if k != "_meta"}


async def _get_api_key_info(request: Request) -> ApiKeyInfo:
    """Extract and validate API key from request."""
    api_key_info = getattr(request.state, "api_key_info", None)
    if api_key_info:
        return api_key_info
    
    # Fallback validation if middleware didn't run
    api_key = request.headers.get("x-api-key") or request.headers.get("authorization", "").replace("Bearer ", "")
    if not api_key:
        raise HTTPException(status_code=401, detail="Missing API key")
    
    redis = await get_redis()
    key_info = await validate_api_key(redis, api_key, get_client_ip(request))
    if not key_info:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    return key_info


async def _get_principal(authorization: Optional[str] = Header(default=None)):
    """Validate JWT bearer token."""
    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")
    token = authorization.replace("Bearer ", "").strip()
    try:
        return verify_token(token, required_scope="voice:session")
    except AuthError as exc:
        raise HTTPException(status_code=401, detail=str(exc)) from exc


# =============================================================================
# Dashboard Routes
# =============================================================================

from app.dashboard.routes import router as dashboard_router
app.include_router(dashboard_router)


# =============================================================================
# Health & Monitoring Endpoints
# =============================================================================

@app.get("/health", response_model=HealthResponse, tags=["Health"])
async def health_check():
    """Health check endpoint for load balancers."""
    global _start_time
    import time
    
    if _start_time is None:
        _start_time = time.time()
    
    return HealthResponse(
        status="healthy",
        version="2.0.0",
        uptime_seconds=time.time() - _start_time,
        latency_stats=get_all_latency_stats(),
    )


@app.get("/ready", tags=["Health"])
async def readiness_check():
    """Readiness check - verifies dependencies are available."""
    try:
        redis = await get_redis()
        await redis.ping()
        return {"status": "ready", "redis": "connected"}
    except Exception as e:
        return JSONResponse(
            status_code=503,
            content={"status": "not_ready", "error": str(e)},
        )


# =============================================================================
# API Key Management Endpoints
# =============================================================================

@app.post("/api/v1/keys", response_model=CreateApiKeyResponse, tags=["API Keys"])
async def create_api_key_endpoint(
    payload: CreateApiKeyRequest,
    request: Request,
    x_admin_key: str = Header(..., description="Admin API key for key management"),
):
    """
    Create a new API key for an organization.
    
    Requires admin API key authentication.
    """
    # Verify admin key
    admin_key = os.getenv("ADMIN_API_KEY")
    if not admin_key or x_admin_key != admin_key:
        raise HTTPException(status_code=403, detail="Invalid admin key")
    
    try:
        plan = Plan(payload.plan)
    except ValueError:
        raise HTTPException(status_code=400, detail=f"Invalid plan: {payload.plan}")
    
    api_key, key_id = generate_api_key(
        org_id=payload.org_id,
        plan=plan,
        allowed_ips=payload.allowed_ips,
        webhook_url=payload.webhook_url,
    )
    
    limits = PLAN_CONFIGS[plan]
    
    return CreateApiKeyResponse(
        api_key=api_key,
        key_id=key_id,
        org_id=payload.org_id,
        plan=plan.value,
        limits={
            "requests_per_minute": limits.requests_per_minute,
            "requests_per_hour": limits.requests_per_hour,
            "requests_per_day": limits.requests_per_day,
            "concurrent_sessions": limits.concurrent_sessions,
            "max_session_duration_seconds": limits.max_session_duration_seconds,
        },
    )


@app.delete("/api/v1/keys", tags=["API Keys"])
async def revoke_api_key_endpoint(
    payload: RevokeApiKeyRequest,
    x_admin_key: str = Header(..., description="Admin API key"),
):
    """Revoke an API key by its ID."""
    admin_key = os.getenv("ADMIN_API_KEY")
    if not admin_key or x_admin_key != admin_key:
        raise HTTPException(status_code=403, detail="Invalid admin key")
    
    redis = await get_redis()
    success = await revoke_api_key(redis, payload.key_id)
    
    if not success:
        raise HTTPException(status_code=500, detail="Failed to revoke key")
    
    return {"status": "revoked", "key_id": payload.key_id}


# =============================================================================
# Authentication Endpoints
# =============================================================================

@app.post("/api/v1/auth/token", tags=["Authentication"])
@timed("auth_token")
async def get_access_token(
    payload: TokenRequest,
    request: Request,
):
    """
    Exchange API key for a short-lived access token.
    
    The access token is used to authenticate subsequent requests.
    """
    key_info = await _get_api_key_info(request)
    
    token, ttl = issue_access_token(
        subject=payload.user_id,
        org=key_info.org_id,
        plan=key_info.plan.value,
        scope=["voice:session", "chat"],
    )
    
    return {
        "access_token": token,
        "token_type": "bearer",
        "expires_in": ttl,
        "org_id": key_info.org_id,
        "plan": key_info.plan.value,
    }


# Legacy endpoint for backward compatibility
@app.post("/api/auth/token", tags=["Authentication"])
async def exchange_api_key(payload: TokenRequest, x_api_key: Optional[str] = Header(default=None)):
    """Legacy token endpoint - use /api/v1/auth/token instead."""
    expected = os.getenv("VOICE_API_KEY")
    if not expected or not x_api_key or x_api_key != expected:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    token, ttl = issue_access_token(
        subject=payload.user_id,
        org=payload.org,
        plan=payload.plan,
        scope=["voice:session"],
    )
    return {"access_token": token, "expires_in": ttl}


# =============================================================================
# Usage & Billing Endpoints
# =============================================================================

@app.get("/api/v1/usage", response_model=UsageResponse, tags=["Usage"])
async def get_usage(request: Request):
    """Get current usage statistics for the organization."""
    key_info = await _get_api_key_info(request)
    redis = await get_redis()
    
    stats = await get_usage_summary(redis, key_info.org_id, key_info.limits)
    
    return UsageResponse(
        org_id=key_info.org_id,
        requests_minute=stats.requests_minute,
        requests_hour=stats.requests_hour,
        requests_day=stats.requests_day,
        concurrent_sessions=stats.concurrent_sessions,
        total_voice_seconds=stats.total_voice_seconds,
        limits={
            "requests_per_minute": key_info.limits.requests_per_minute,
            "requests_per_hour": key_info.limits.requests_per_hour,
            "requests_per_day": key_info.limits.requests_per_day,
            "concurrent_sessions": key_info.limits.concurrent_sessions,
        },
    )


@app.get("/api/v1/limits", tags=["Usage"])
async def get_rate_limits(request: Request):
    """Get current rate limit headers without counting as a request."""
    key_info = await _get_api_key_info(request)
    redis = await get_redis()
    
    headers = await get_remaining_quota(redis, key_info.org_id, key_info.limits)
    return headers


# =============================================================================
# Voice Session Endpoints
# =============================================================================

@app.post("/api/v1/voice/session", response_model=VoiceSessionResponse, tags=["Voice"])
@timed("voice_session_create")
async def create_voice_session_v1(
    request: Request,
    payload: VoiceSessionRequest,
):
    """
    Create a new voice session.
    
    Returns a WebSocket URL and token for connecting to the voice stream.
    """
    key_info = await _get_api_key_info(request)
    redis = await get_redis()
    
    # Reserve session slot
    sid = str(uuid.uuid4())
    try:
        await reserve_session(redis, key_info.org_id, sid, key_info.limits)
    except ConcurrencyLimitError as e:
        raise HTTPException(
            status_code=429,
            detail=f"Concurrent session limit reached: {e.current}/{e.limit}",
        )
    
    # Issue voice token
    voice_token, ttl = issue_voice_token(
        subject=f"org:{key_info.org_id}",
        org=key_info.org_id,
        plan=key_info.plan.value,
        scope=["voice:connect"],
        sid=sid,
    )
    
    ws_url = _ws_url_from_request(request)
    
    _logger.info(
        "voice_session_created",
        org_id=key_info.org_id,
        sid=sid,
        plan=key_info.plan.value,
    )
    
    return VoiceSessionResponse(
        sid=sid,
        token=voice_token,
        ws_url=ws_url,
        expires_in=ttl,
    )


# Legacy endpoint
@app.post("/api/voice/session", response_model=VoiceSessionResponse, tags=["Voice"])
async def create_voice_session(
    request: Request,
    payload: VoiceSessionRequest,
    principal=Depends(_get_principal),
):
    """Legacy voice session endpoint."""
    sid = str(uuid.uuid4())
    plan = principal.get("plan", "free")
    
    voice_token, ttl = issue_voice_token(
        subject=principal["sub"],
        org=principal.get("org"),
        plan=plan,
        scope=["voice:connect"],
        sid=sid,
    )
    
    ws_url = _ws_url_from_request(request)
    return VoiceSessionResponse(sid=sid, token=voice_token, ws_url=ws_url, expires_in=ttl)


@app.websocket("/voice")
async def voice_ws(websocket: WebSocket):
    """WebSocket endpoint for voice streaming."""
    await voice_ws_endpoint(websocket, get_redis)


# =============================================================================
# Chat Endpoints
# =============================================================================

@app.post("/api/v1/chat", response_model=ChatResponse, tags=["Chat"])
@timed("chat")
async def chat_v1(payload: ChatRequest, request: Request) -> ChatResponse:
    """
    Send a chat message and receive AI response.
    
    Authenticated via API key.
    """
    key_info = await _get_api_key_info(request)
    return await _process_chat(payload, key_info.org_id)


@app.post("/api/chat", response_model=ChatResponse, tags=["Chat"])
async def chat(payload: ChatRequest) -> ChatResponse:
    """Legacy chat endpoint."""
    return await _process_chat(payload, None)


async def _process_chat(payload: ChatRequest, org_id: Optional[str]) -> ChatResponse:
    """Process chat message."""
    message = payload.message.strip()
    if not message:
        raise HTTPException(status_code=400, detail="Message is required")
    
    session_id = payload.session_id or str(uuid.uuid4())
    session = _chat_sessions.get(session_id)
    if session is None:
        session = ChatSession()
        _chat_sessions[session_id] = session
    
    tracker = get_latency_tracker("llm_call")
    import time
    start = time.perf_counter()
    
    try:
        output = _chat_brain.run_turn(session.sales_state, message, session.memory)
    except APIConnectionError as exc:
        raise HTTPException(
            status_code=503,
            detail="LLM connection error. Check AZURE_OPENAI_ENDPOINT and network/DNS.",
        ) from exc
    except APIError as exc:
        raise HTTPException(status_code=502, detail=f"LLM API error: {exc}") from exc
    finally:
        tracker.record((time.perf_counter() - start) * 1000)
    
    session.sales_state = output.next_state
    session.turns += 1
    merge_memory(session.memory, output)
    
    return ChatResponse(
        session_id=session_id,
        reply=output.reply,
        state=output.next_state,
        lead_score=output.lead_score,
        missing_fields=output.missing_fields,
        memory=_strip_meta(session.memory),
        safety_flags=output.safety_flags,
        turns=session.turns,
    )


# =============================================================================
# Static Files & Web UI
# =============================================================================

@app.get("/", tags=["Web UI"])
async def webui_index() -> FileResponse:
    """Serve the web UI."""
    return FileResponse(os.path.join(WEB_ROOT, "index.html"))


# =============================================================================
# Error Handlers
# =============================================================================

@app.exception_handler(RateLimitError)
async def rate_limit_error_handler(request: Request, exc: RateLimitError):
    """Handle rate limit errors."""
    return JSONResponse(
        status_code=429,
        content={
            "error": str(exc),
            "limit_type": exc.limit_type,
            "limit": exc.limit,
            "current": exc.current,
            "retry_after": exc.retry_after,
        },
        headers={"Retry-After": str(exc.retry_after)},
    )


@app.exception_handler(ConcurrencyLimitError)
async def concurrency_limit_error_handler(request: Request, exc: ConcurrencyLimitError):
    """Handle concurrent session limit errors."""
    return JSONResponse(
        status_code=429,
        content={
            "error": str(exc),
            "limit": exc.limit,
            "current": exc.current,
        },
    )


# For backward compatibility with existing QuotaError
from app.auth.rate_limit import QuotaError

@app.exception_handler(QuotaError)
async def quota_error_handler(_: Request, exc: QuotaError):
    """Handle legacy quota errors."""
    return JSONResponse(status_code=429, content={"detail": str(exc)})
