use anyhow::Result;
use axum::{
    body::Body,
    extract::State,
    http::{HeaderMap, Request, StatusCode},
    response::{IntoResponse, Response, Sse},
    routing::post,
    Router,
};
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{info, warn};

use super::middleware;

use crate::auth::AuthProvider;
use crate::config::AppConfig;
use crate::data::{EventQueue, TelemetryEvent};
use crate::router::RouterBrain;

pub struct ProxyState {
    pub auth: Arc<dyn AuthProvider>,
    pub router: Arc<RouterBrain>,
    pub telemetry_tx: mpsc::Sender<TelemetryEvent>,
    pub http_client: reqwest::Client,
    pub config: Arc<AppConfig>,
}

pub fn create_router(state: Arc<ProxyState>) -> Router {
    Router::new()
        .route("/v1/chat/completions", post(proxy_chat_completions))
        .route("/v1/completions", post(proxy_completions))
        .route("/v1/embeddings", post(proxy_embeddings))
        .route("/health", axum::routing::get(health))
        .with_state(state)
}

async fn health() -> &'static str {
    "ok"
}

async fn proxy_chat_completions(
    State(state): State<Arc<ProxyState>>,
    headers: HeaderMap,
    body: String,
) -> Result<Response, ProxyError> {
    // 1. Authenticate
    let auth_ctx = state.auth.authenticate(&headers).await?;

    // 2. Parse request
    let mut request: serde_json::Value =
        serde_json::from_str(&body).map_err(|_| ProxyError::BadRequest("Invalid JSON"))?;

    let is_streaming = request
        .get("stream")
        .and_then(|v| v.as_bool())
        .unwrap_or(false);

    let original_model = request
        .get("model")
        .and_then(|v| v.as_str())
        .unwrap_or("gpt-4o")
        .to_string();

    // 3. Route (pick model + provider)
    let decision = state
        .router
        .route(&auth_ctx.org_id, &request)
        .await;

    // Apply routing decision
    if let Some(ref routed_model) = decision.model {
        request["model"] = serde_json::Value::String(routed_model.clone());
    }

    let provider = decision.provider.unwrap_or_default();
    let upstream_url = state.config.provider_url(&provider);

    // 4. Forward to upstream
    let start = std::time::Instant::now();

    let upstream_resp = state
        .http_client
        .post(format!("{}/v1/chat/completions", upstream_url))
        .header("Authorization", format!("Bearer {}", state.config.provider_key(&provider)))
        .header("Content-Type", "application/json")
        .body(request.to_string())
        .send()
        .await
        .map_err(|e| ProxyError::UpstreamError(e.to_string()))?;

    let latency = start.elapsed();
    let status = upstream_resp.status();

    // 5. Emit telemetry (non-blocking)
    let _ = state.telemetry_tx.try_send(TelemetryEvent {
        org_id: auth_ctx.org_id.clone(),
        original_model: original_model.clone(),
        routed_model: decision.model.clone().unwrap_or(original_model),
        provider: provider.clone(),
        strategy: decision.strategy.clone(),
        latency_ms: latency.as_millis() as u32,
        status_code: status.as_u16(),
        is_streaming,
        prompt_tokens: 0,  // Filled by worker from response
        completion_tokens: 0,
        timestamp: chrono::Utc::now(),
    });

    // 6. Pass through response transparently
    let resp_headers = upstream_resp.headers().clone();
    let resp_status = upstream_resp.status();

    if is_streaming {
        // Stream SSE chunks directly
        let byte_stream = upstream_resp.bytes_stream();
        let body = Body::from_stream(byte_stream);

        let mut response = Response::builder().status(resp_status);
        for (key, value) in resp_headers.iter() {
            response = response.header(key, value);
        }
        Ok(response.body(body).unwrap())
    } else {
        let resp_body = upstream_resp
            .bytes()
            .await
            .map_err(|e| ProxyError::UpstreamError(e.to_string()))?;

        let mut response = Response::builder().status(resp_status);
        for (key, value) in resp_headers.iter() {
            response = response.header(key, value);
        }
        Ok(response.body(Body::from(resp_body)).unwrap())
    }
}

// Placeholder — same pattern as chat completions
async fn proxy_completions(
    State(state): State<Arc<ProxyState>>,
    headers: HeaderMap,
    body: String,
) -> Result<Response, ProxyError> {
    proxy_chat_completions(State(state), headers, body).await
}

async fn proxy_embeddings(
    State(state): State<Arc<ProxyState>>,
    headers: HeaderMap,
    body: String,
) -> Result<Response, ProxyError> {
    // Embeddings don't need routing — pass through directly
    let auth_ctx = state.auth.authenticate(&headers).await?;
    let provider = "openai".to_string();
    let upstream_url = state.config.provider_url(&provider);

    let resp = state
        .http_client
        .post(format!("{}/v1/embeddings", upstream_url))
        .header("Authorization", format!("Bearer {}", state.config.provider_key(&provider)))
        .header("Content-Type", "application/json")
        .body(body)
        .send()
        .await
        .map_err(|e| ProxyError::UpstreamError(e.to_string()))?;

    let status = resp.status();
    let body = resp.bytes().await.map_err(|e| ProxyError::UpstreamError(e.to_string()))?;
    Ok(Response::builder().status(status).body(Body::from(body)).unwrap())
}

// --- Error types ---

#[derive(Debug, thiserror::Error)]
pub enum ProxyError {
    #[error("Authentication failed: {0}")]
    AuthError(String),
    #[error("Bad request: {0}")]
    BadRequest(&'static str),
    #[error("Upstream error: {0}")]
    UpstreamError(String),
}

impl From<crate::auth::AuthError> for ProxyError {
    fn from(e: crate::auth::AuthError) -> Self {
        ProxyError::AuthError(e.to_string())
    }
}

impl IntoResponse for ProxyError {
    fn into_response(self) -> Response {
        let (status, msg) = match &self {
            ProxyError::AuthError(_) => (StatusCode::UNAUTHORIZED, self.to_string()),
            ProxyError::BadRequest(_) => (StatusCode::BAD_REQUEST, self.to_string()),
            ProxyError::UpstreamError(_) => (StatusCode::BAD_GATEWAY, self.to_string()),
        };
        (status, serde_json::json!({"error": msg}).to_string()).into_response()
    }
}
