from __future__ import annotations

import os
import time
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, Response, 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.rate_limit import QuotaError
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,
    reserve_session,
    release_session,
    get_usage_summary,
    get_remaining_quota,
    check_rate_limit,
)
from app.auth.security import SecurityMiddleware, get_client_ip
from app.storage.redis import get_redis, close_redis
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
from app.billing import (
    init_billing,
    shutdown_billing,
    track_voice_session_start,
    track_api_request,
    track_rate_limit_hit,
    get_usage_for_period,
    calculate_invoice_amount,
)
from app.dashboard.routes import router as dashboard_router

if load_dotenv:
    load_dotenv()

configure_logging()
_logger = get_logger("main")


# Lifespan management for startup/shutdown
@asynccontextmanager
async def lifespan(app: FastAPI):
    _logger.info("starting_up")
    await warm_up_connections(get_redis)
    await init_billing()
    
    # Initialize dashboard demo data
    from app.dashboard.data import init_demo_data
    await init_demo_data()
    
    yield
    _logger.info("shutting_down")
    await shutdown_billing()
    await close_redis()


app = FastAPI(
    title="SalesAI API",
    description="""## AI-Powered Sales Assistant API

SalesAI provides intelligent conversation management, lead scoring, and voice capabilities for sales teams.

### Key Features
- **Chat API**: AI-powered sales conversation responses
- **Voice API**: Real-time voice sessions with STT/TTS
- **Analytics**: Lead scoring, sentiment analysis, and sales insights
- **Multi-tenant**: Organization-based API keys with usage tracking

### Authentication
All endpoints require authentication via API key (`X-API-Key` header) or Bearer token.

### Rate Limits
- **Free**: 100 requests/day
- **Starter**: 1,000 requests/month
- **Pro**: 10,000 requests/month  
- **Enterprise**: 100,000 requests/month
""",
    version="2.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json",
    openapi_tags=[
        {"name": "Authentication", "description": "API key and token management"},
        {"name": "Chat", "description": "AI-powered sales conversations"},
        {"name": "Voice", "description": "Real-time voice sessions"},
        {"name": "Usage", "description": "Usage statistics and billing"},
        {"name": "API Keys", "description": "API key management (admin only)"},
        {"name": "Health", "description": "Health and readiness checks"},
    ],
    contact={
        "name": "SalesAI Support",
        "email": "support@salesai.com",
    },
    license_info={
        "name": "Proprietary",
    },
)

# Middleware
app.add_middleware(GZipMiddleware, minimum_size=1000)
app.add_middleware(
    CORSMiddleware,
    allow_origins=os.getenv("CORS_ORIGINS", "*").split(","),
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["X-Request-ID", "X-Response-Time", "X-RateLimit-Remaining-Minute", "Retry-After"],
)
app.add_middleware(SecurityMiddleware, redis_factory=get_redis)

# Include routers
app.include_router(dashboard_router)

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


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 ApiKeyTokenRequest(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


PLAN_CONCURRENCY = {
    "free": 1,
    "pro": 3,
    "enterprise": 20,
}


@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()


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_principal(authorization: Optional[str] = Header(default=None)):
    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


@app.post("/api/auth/token", tags=["Authentication"])
async def exchange_api_key(payload: ApiKeyTokenRequest, x_api_key: Optional[str] = Header(default=None)):
    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}


@app.get("/")
async def webui_index() -> FileResponse:
    return FileResponse(os.path.join(WEB_ROOT, "index.html"))


@app.post("/api/chat", response_model=ChatResponse, tags=["Chat"], summary="Generate AI sales response", description="Send a message and receive an AI-powered sales response with lead scoring and context extraction.")
async def chat(
    payload: ChatRequest,
    request: Request,
    response: Response,
    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
) -> ChatResponse:
    message = payload.message.strip()
    if not message:
        raise HTTPException(status_code=400, detail="Message is required")

    require_api_key = os.getenv("REQUIRE_API_KEY", "true").lower() in ("1", "true", "yes")
    api_key_info: Optional[ApiKeyInfo] = None
    redis = await get_redis()
    client_ip = get_client_ip(request)

    if x_api_key:
        api_key_info = await validate_api_key(redis, x_api_key, client_ip=client_ip)
        if not api_key_info:
            raise HTTPException(status_code=401, detail="Invalid API key")
    elif require_api_key:
        raise HTTPException(status_code=401, detail="Missing API key")

    start_time = time.perf_counter()
    status_code = 200

    if api_key_info:
        try:
            await check_rate_limit(redis, api_key_info.org_id, api_key_info.limits, endpoint="/api/chat")
            remaining = await get_remaining_quota(redis, api_key_info.org_id, api_key_info.limits)
            response.headers.update({k: str(v) for k, v in remaining.items()})
        except RateLimitError as exc:
            await track_rate_limit_hit(
                api_key_info.org_id,
                exc.limit_type,
                exc.limit,
                exc.current,
            )
            status_code = 429
            raise HTTPException(
                status_code=429,
                detail=str(exc),
                headers={"Retry-After": str(exc.retry_after)},
            ) from exc

    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
    try:
        output = _chat_brain.run_turn(session.sales_state, message, session.memory)
    except APIConnectionError as exc:
        status_code = 503
        raise HTTPException(
            status_code=503,
            detail="LLM connection error. Check AZURE_OPENAI_ENDPOINT and network/DNS.",
        ) from exc
    except APIError as exc:
        status_code = 502
        raise HTTPException(status_code=502, detail=f"LLM API error: {exc}") from exc
    finally:
        if api_key_info:
            duration_ms = int((time.perf_counter() - start_time) * 1000)
            await track_api_request(
                api_key_info.org_id,
                endpoint="/api/chat",
                method=request.method,
                status_code=status_code,
                duration_ms=duration_ms,
            )

    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,
    )


@app.post("/api/voice/session", response_model=VoiceSessionResponse, tags=["Voice"])
async def create_voice_session(request: Request, payload: VoiceSessionRequest, principal=Depends(_get_principal)):
    sid = str(uuid.uuid4())
    plan = principal.get("plan", "free")
    scope = ["voice:connect"]
    voice_token, ttl = issue_voice_token(
        subject=principal["sub"],
        org=principal.get("org"),
        plan=plan,
        scope=scope,
        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):
    await voice_ws_endpoint(websocket, get_redis)


# =============================================================================
# Production API v1 Endpoints
# =============================================================================

class CreateApiKeyRequest(BaseModel):
    org_id: str = Field(..., description="Organization identifier")
    plan: str = Field(default="free", description="Pricing plan: free, starter, pro, enterprise")
    allowed_ips: Optional[List[str]] = Field(default=None, description="IP allowlist")


class CreateApiKeyResponse(BaseModel):
    api_key: str
    key_id: str
    org_id: str
    plan: str
    limits: Dict[str, Any]


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]


async def _get_api_key_info(request: Request) -> ApiKeyInfo:
    """Extract and validate API key from request state or headers."""
    api_key_info = getattr(request.state, "api_key_info", None)
    if api_key_info:
        return api_key_info
    
    api_key = request.headers.get("x-api-key") or request.headers.get("authorization", "").replace("Bearer ", "")
    if not api_key or not api_key.startswith("sv_"):
        raise HTTPException(status_code=401, detail="Missing or invalid 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


@app.get("/health", tags=["Health"])
async def health_check():
    """Health check for load balancers."""
    return {
        "status": "healthy",
        "version": "2.0.0",
        "latency_stats": get_all_latency_stats(),
    }


@app.get("/ready", tags=["Health"])
async def readiness_check():
    """Readiness check - verifies dependencies."""
    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)})


@app.post("/api/v1/keys", response_model=CreateApiKeyResponse, tags=["API Keys"])
async def create_api_key_endpoint(
    payload: CreateApiKeyRequest,
    x_admin_key: str = Header(..., description="Admin API key"),
):
    """Create a new API key for an organization (requires 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,
    )
    
    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.post("/api/v1/auth/token", tags=["Authentication"])
@timed("auth_token")
async def get_access_token_v1(payload: ApiKeyTokenRequest, request: Request):
    """Exchange API key for access token."""
    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,
    }


@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 voice session with API key auth and rate limiting."""
    key_info = await _get_api_key_info(request)
    redis = await get_redis()
    
    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: {e.current}/{e.limit}")
    
    # Track for billing
    await track_voice_session_start(key_info.org_id, sid, key_info.plan.value)
    
    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,
    )
    
    _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_from_request(request),
        expires_in=ttl,
    )


@app.get("/api/v1/usage", response_model=UsageResponse, tags=["Usage"])
async def get_usage(request: Request):
    """Get current usage statistics."""
    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/usage/billing", tags=["Usage"])
async def get_billing_usage(
    request: Request,
    start: int = 0,
    end: int = 0,
):
    """Get usage for a billing period (timestamps in seconds)."""
    import time
    key_info = await _get_api_key_info(request)
    redis = await get_redis()
    
    # Default to current month
    if start == 0:
        now = time.time()
        start = int(now - (now % 86400) - 30 * 86400)  # 30 days ago
    if end == 0:
        end = int(time.time())
    
    usage = await get_usage_for_period(redis, key_info.org_id, start, end)
    invoice = calculate_invoice_amount(usage, key_info.plan.value)
    
    return {
        "usage": usage,
        "invoice_preview": invoice,
    }


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

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


@app.exception_handler(RateLimitError)
async def rate_limit_error_handler(_: Request, exc: RateLimitError):
    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, exc: ConcurrencyLimitError):
    return JSONResponse(
        status_code=429,
        content={"error": str(exc), "limit": exc.limit, "current": exc.current},
    )
