use actix_files as fs;
use actix_web::{web, App, HttpServer, Responder, HttpResponse, Result};
use serde::{Deserialize, Serialize};
use tera::{Tera, Context};
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc;
use tokio_tungstenite::{connect_async, tungstenite};
use futures_util::{StreamExt, SinkExt};
use secp256k1::{Secp256k1, SecretKey, Message};
use sha2::{Sha256, Digest};
use url::Url;
use std::thread;

// Struct for representing a sign request
#[derive(Serialize, Deserialize, Clone, Debug)]
struct SignRequest {
    id: String,
    message: String,
    #[serde(skip)]
    message_raw: String, // Original base64 message for sending back in the response
    #[serde(skip)]
    message_decoded: String, // Decoded message for display
}

// Struct for representing the application state
struct AppState {
    templates: Tera,
    keypair: Arc<KeyPair>,
    pending_request: Arc<Mutex<Option<SignRequest>>>,
    websocket_sender: mpsc::Sender<WebSocketCommand>,
}

// Commands that can be sent to the WebSocket connection
enum WebSocketCommand {
    Sign { id: String, message: String, signature: Vec<u8> },
    Close,
}

// Keypair for signing messages
struct KeyPair {
    secret_key: SecretKey,
    public_key_hex: String,
}

impl KeyPair {
    fn new() -> Self {
        let secp = Secp256k1::new();
        let mut rng = rand::thread_rng();
        
        // Generate a new random keypair
        let (secret_key, public_key) = secp.generate_keypair(&mut rng);
        
        // Convert public key to hex for identification
        let public_key_hex = hex::encode(public_key.serialize());
        
        KeyPair {
            secret_key,
            public_key_hex,
        }
    }
    
    fn sign(&self, message: &[u8]) -> Vec<u8> {
        // Hash the message first (secp256k1 requires a 32-byte hash)
        let mut hasher = Sha256::new();
        hasher.update(message);
        let message_hash = hasher.finalize();
        
        // Create a secp256k1 message from the hash
        let secp_message = Message::from_slice(&message_hash).unwrap();
        
        // Sign the message
        let secp = Secp256k1::new();
        let signature = secp.sign_ecdsa(&secp_message, &self.secret_key);
        
        // Return the serialized signature
        signature.serialize_compact().to_vec()
    }
}

// Controller for the index page
async fn index(data: web::Data<AppState>) -> Result<HttpResponse> {
    let mut context = Context::new();
    
    // Add the keypair to the context
    context.insert("public_key", &data.keypair.public_key_hex);
    
    // Add the pending request if there is one
    if let Some(request) = &*data.pending_request.lock().unwrap() {
        context.insert("request", request);
    }
    
    let rendered = data.templates.render("index.html", &context)
        .map_err(|e| {
            eprintln!("Template error: {}", e);
            actix_web::error::ErrorInternalServerError("Template error")
        })?;
    
    Ok(HttpResponse::Ok().content_type("text/html").body(rendered))
}

// Controller for the sign endpoint
async fn sign_request(
    data: web::Data<AppState>,
    form: web::Form<SignRequestForm>,
) -> impl Responder {
    println!("SIGN ENDPOINT: Starting sign_request handler for form ID: {}", form.id);
    
    // Try to get a lock on the pending request
    println!("SIGN ENDPOINT: Attempting to acquire lock on pending_request");
    match data.pending_request.try_lock() {
        Ok(mut guard) => {
            // Check if we have a pending request
            if let Some(request) = &*guard {
                println!("SIGN ENDPOINT: Found pending request with ID: {}", request.id);
                
                // Get the request ID
                let id = request.id.clone();
                
                // Verify that the request ID matches
                if id == form.id {
                    println!("SIGN ENDPOINT: Request ID matches form ID: {}", id);
                    
                    // Sign the message
                    let message = request.message.as_bytes();
                    println!("SIGN ENDPOINT: About to sign message: {} (length: {})", 
                          String::from_utf8_lossy(message), message.len());
                    let signature = data.keypair.sign(message);
                    println!("SIGN ENDPOINT: Message signed successfully. Signature length: {}", signature.len());
                    
                    // Send the signature via WebSocket
                    println!("SIGN ENDPOINT: About to send signature via websocket channel");
                    match data.websocket_sender.send(WebSocketCommand::Sign { 
                        id: id.clone(), 
                        message: request.message_raw.clone(), // Include the original base64 message
                        signature 
                    }).await {
                        Ok(_) => {
                            println!("SIGN ENDPOINT: Successfully sent signature to websocket channel");
                        },
                        Err(e) => {
                            let error_msg = format!("Failed to send signature: {}", e);
                            println!("SIGN ENDPOINT ERROR: {}", error_msg);
                            return HttpResponse::InternalServerError()
                                .content_type("text/html")
                                .body(format!("<h1>Error sending signature</h1><p>{}</p><p><a href='/'>Return to home</a></p>", error_msg));
                        }
                    }
                    
                    // Clear the pending request
                    println!("SIGN ENDPOINT: Clearing pending request");
                    *guard = None;
                    
                    // Return a success page that continues to the next step
                    println!("SIGN ENDPOINT: Returning success response");
                    return HttpResponse::Ok()
                        .content_type("text/html")
                        .body(r#"<html>
                            <head>
                                <title>Signature Sent</title>
                                <meta http-equiv="refresh" content="2; url=/" />
                                <script type="text/javascript">
                                    console.log("Signature sent successfully, redirecting in 2 seconds...");
                                    setTimeout(function() { window.location.href = '/'; }, 2000);
                                </script>
                                <style>
                                    body { font-family: Arial, sans-serif; text-align: center; margin-top: 50px; }
                                    .success { color: green; }
                                </style>
                            </head>
                            <body>
                                <h1 class="success">✓ Signature Sent Successfully!</h1>
                                <p>Redirecting back to home page...</p>
                                <p><a href="/">Click here if you're not redirected automatically</a></p>
                            </body>
                        </html>"#);
                } else {
                    println!("SIGN ENDPOINT: Request ID {} does not match form ID {}", request.id, form.id);
                }
            } else {
                println!("SIGN ENDPOINT: No pending request found");
            }
        },
        Err(e) => {
            let error_msg = format!("Failed to acquire lock on pending_request: {}", e);
            println!("SIGN ENDPOINT ERROR: {}", error_msg);
            return HttpResponse::InternalServerError()
                .content_type("text/html")
                .body(format!("<h1>Error processing request</h1><p>{}</p><p><a href='/'>Return to home</a></p>", error_msg));
        }
    }
    
    // Redirect back to the index page (if no request was found or ID didn't match)
    println!("SIGN ENDPOINT: No matching request found, redirecting to home");
    HttpResponse::SeeOther()
        .append_header(("Location", "/"))
        .finish()
}

// Form for submitting a signature
#[derive(Deserialize)]
struct SignRequestForm {
    id: String,
}

// WebSocket client task that connects to the SigSocket server
async fn websocket_client_task(
    keypair: Arc<KeyPair>, 
    pending_request: Arc<Mutex<Option<SignRequest>>>,
    mut command_receiver: mpsc::Receiver<WebSocketCommand>,
) {
    // Connect directly to the web app's integrated SigSocket endpoint
    let sigsocket_url = "ws://127.0.0.1:8080/ws";
    
    // Reconnection settings
    let mut retry_count = 0;
    const MAX_RETRY_COUNT: u32 = 10; // Reset retry counter after this many attempts
    const BASE_RETRY_DELAY_MS: u64 = 1000; // Start with 1 second
    const MAX_RETRY_DELAY_MS: u64 = 30000; // Cap at 30 seconds
    
    loop {
        // Calculate backoff delay with jitter for retry
        let delay_ms = if retry_count > 0 {
            let base_delay = BASE_RETRY_DELAY_MS * 2u64.pow(retry_count.min(6));
            let jitter = rand::random::<u64>() % 500; // Add up to 500ms of jitter
            (base_delay + jitter).min(MAX_RETRY_DELAY_MS)
        } else {
            0 // No delay on first attempt
        };
        
        if retry_count > 0 {
            println!("Reconnection attempt {} in {} ms...", retry_count, delay_ms);
            tokio::time::sleep(tokio::time::Duration::from_millis(delay_ms)).await;
        }
        
        // Connect to the SigSocket server with timeout
        println!("Connecting to SigSocket server at {}", sigsocket_url);
        let connect_result = tokio::time::timeout(
            tokio::time::Duration::from_secs(10), // Connection timeout
            connect_async(Url::parse(sigsocket_url).unwrap())
        ).await;
        
        match connect_result {
            // Timeout error
            Err(_) => {
                eprintln!("Connection attempt timed out");
                retry_count = (retry_count + 1) % MAX_RETRY_COUNT;
                continue;
            },
            // Connection result
            Ok(conn_result) => match conn_result {
                // Connection successful
                Ok((mut ws_stream, _)) => {
                    println!("Connected to SigSocket server");
                    // Reset retry counter on successful connection
                    retry_count = 0;
                    
                    // Heartbeat functionality has been removed
                    println!("DEBUG: Running without heartbeat functionality");
                    
                    // Send the initial message with just the raw public key
                    let intro_message = keypair.public_key_hex.clone();
                    if let Err(e) = ws_stream.send(tungstenite::Message::Text(intro_message)).await {
                        eprintln!("Failed to send introduction message: {}", e);
                        continue;
                    }
                    
                    println!("Sent introduction with public key: {}", keypair.public_key_hex);
                    
                    // Last time we received a message or pong from the server
                    let mut last_server_response = std::time::Instant::now();
                    
                    // Process incoming messages and commands
                    loop {
                        tokio::select! {
                            // Handle WebSocket message
                            msg = ws_stream.next() => {
                                match msg {
                                    Some(Ok(tungstenite::Message::Text(text))) => {
                                        println!("Received message: {}", text);
                                        last_server_response = std::time::Instant::now();
                                        
                                        // Parse the message as a sign request
                                        match serde_json::from_str::<SignRequest>(&text) {
                                            Ok(mut request) => {
                                                println!("DEBUG: Successfully parsed sign request with ID: {}", request.id);
                                                println!("DEBUG: Base64 message: {}", request.message);
                                                
                                                // Save the original base64 message for later use in response
                                                request.message_raw = request.message.clone();
                                                
                                                // Decode the base64 message content
                                                match base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &request.message) {
                                                    Ok(decoded) => {
                                                        let decoded_text = String::from_utf8_lossy(&decoded).to_string();
                                                        println!("DEBUG: Decoded message: {}", decoded_text);
                                                        
                                                        // Store the decoded message for display
                                                        request.message_decoded = decoded_text;
                                                        
                                                        // Update the message for displaying in the UI
                                                        request.message = request.message_decoded.clone();
                                                        
                                                        // Store the request for display in the UI
                                                        *pending_request.lock().unwrap() = Some(request);
                                                        println!("Received signing request. Please check the web UI to approve it.");
                                                    },
                                                    Err(e) => {
                                                        eprintln!("Error decoding base64 message: {}", e);
                                                    }
                                                }
                                            },
                                            Err(e) => {
                                                eprintln!("Error parsing sign request JSON: {}", e);
                                                eprintln!("Raw message: {}", text);
                                            }
                                        }
                                    },
                                    Some(Ok(tungstenite::Message::Ping(data))) => {
                                        // Respond to ping with pong
                                        last_server_response = std::time::Instant::now();
                                        if let Err(e) = ws_stream.send(tungstenite::Message::Pong(data)).await {
                                            eprintln!("Failed to send pong: {}", e);
                                            break;
                                        }
                                    },
                                    Some(Ok(tungstenite::Message::Pong(_))) => {
                                        // Got pong response from the server
                                        last_server_response = std::time::Instant::now();
                                    },
                                    Some(Ok(_)) => {
                                        // Ignore other types of messages
                                        last_server_response = std::time::Instant::now();
                                    },
                                    Some(Err(e)) => {
                                        eprintln!("WebSocket error: {}", e);
                                        break;
                                    },
                                    None => {
                                        eprintln!("WebSocket connection closed");
                                        break;
                                    },
                                }
                            },
                            
                            // Heartbeat functionality has been removed
                            
                            // Handle signing command from the web interface
                            cmd = command_receiver.recv() => {
                                match cmd {
                                    Some(WebSocketCommand::Sign { id, message, signature }) => {
                                        println!("DEBUG: Signing request ID: {}", id);
                                        println!("DEBUG: Raw signature bytes: {:?}", signature);
                                        println!("DEBUG: Using message from command: {}", message);
                                        
                                        // Convert signature bytes to base64
                                        let sig_base64 = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &signature);
                                        println!("DEBUG: Base64 signature: {}", sig_base64);
                                        
                                        // Create a JSON response with explicit ID and message/signature fields
                                        let response = format!("{{\"id\": \"{}\", \"message\": \"{}\", \"signature\": \"{}\"}}", 
                                                             id, message, sig_base64);
                                        println!("DEBUG: Preparing to send JSON response: {}", response);
                                        println!("DEBUG: Response length: {} bytes", response.len());
                                        
                                        // Log that we're about to send on the WebSocket connection
                                        println!("DEBUG: About to send on WebSocket connection");
                                        
                                        // Send the signature response right away - with extra logging
                                        println!("!!!! ATTEMPTING TO SEND SIGNATURE RESPONSE NOW !!!!");
                                        match ws_stream.send(tungstenite::Message::Text(response.clone())).await {
                                            Ok(_) => {
                                                last_server_response = std::time::Instant::now();
                                                println!("!!!! SUCCESSFULLY SENT SIGNATURE RESPONSE !!!!");
                                                println!("!!!! SIGNATURE SENT FOR REQUEST ID: {} !!!!", id);
                                                
                                                // Clear the pending request after successful signature
                                                *pending_request.lock().unwrap() = None;
                                                
                                                // Send another simple message to confirm the connection is still working
                                                if let Err(e) = ws_stream.send(tungstenite::Message::Text("CONFIRM_SIGNATURE_SENT".to_string())).await {
                                                    println!("DEBUG: Failed to send confirmation message: {}", e);
                                                } else {
                                                    println!("DEBUG: Sent confirmation message after signature");
                                                }
                                            },
                                            Err(e) => {
                                                eprintln!("!!!! FAILED TO SEND SIGNATURE RESPONSE: {} !!!!", e);
                                                // Try to reconnect or recover
                                                println!("DEBUG: Attempting to diagnose connection issue...");
                                                break;
                                            }
                                        }
                                    },
                                    Some(WebSocketCommand::Close) => {
                                        println!("DEBUG: Received close command, closing connection");
                                        break;
                                    },
                                    None => {
                                        eprintln!("Command channel closed");
                                        break;
                                    }
                                }
                            }
                        }
                    }
                    
                    // Connection loop has ended, will attempt to reconnect
                    println!("WebSocket connection closed, will attempt to reconnect...");
                },
                
                // Connection error
                Err(e) => {
                    eprintln!("Failed to connect to SigSocket server: {}", e);
                }
            }
        }
        
        // Increment retry counter but don't exceed MAX_RETRY_COUNT
        retry_count = (retry_count + 1) % MAX_RETRY_COUNT;
    }
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    // Setup logger
    env_logger::init_from_env(env_logger::Env::default().default_filter_or("info"));
    
    // Initialize templates
    let mut tera = Tera::default();
    tera.add_raw_templates(vec![
        ("index.html", include_str!("../templates/index.html")),
    ]).unwrap();
    
    // Generate a keypair for signing
    let keypair = Arc::new(KeyPair::new());
    println!("Generated keypair with public key: {}", keypair.public_key_hex);
    
    // Create a channel for sending commands to the WebSocket client
    let (command_sender, command_receiver) = mpsc::channel::<WebSocketCommand>(32);
    
    // Create the pending request mutex
    let pending_request = Arc::new(Mutex::new(None::<SignRequest>));
    
    // Spawn the WebSocket client task
    let ws_keypair = keypair.clone();
    let ws_pending_request = pending_request.clone();
    tokio::spawn(async move {
        websocket_client_task(ws_keypair, ws_pending_request, command_receiver).await;
    });
    
    // Create the app state
    let app_state = web::Data::new(AppState {
        templates: tera,
        keypair,
        pending_request,
        websocket_sender: command_sender,
    });
    
    println!("Client App server starting on http://127.0.0.1:8082");
    
    // Start the web server
    HttpServer::new(move || {
        App::new()
            .app_data(app_state.clone())
            // Register routes
            .route("/", web::get().to(index))
            .route("/sign", web::post().to(sign_request))
            // Static files
            .service(fs::Files::new("/static", "./static"))
    })
    .bind("127.0.0.1:8082")?
    .run()
    .await
}