//! The channel server: one turn per connection over the M1 frame protocol. A session runs one
//! turn at a time; the busy set keeps a second turn on the same session from starting until the
//! first finishes, and is released before every final frame so a client can send the next turn the
//! moment it reads the last one.

use std::collections::HashSet;
use std::os::unix::net::{UnixListener, UnixStream};
use std::sync::{Arc, Mutex};

use crate::baseline::Baseline;
use crate::config::Config;
use crate::llama::Client;
use crate::session::{Session, SessionError};
use crate::tools::{Registry, ToolPort};
use crate::turn::{Runtime, TurnError, run_turn};
use proto::{
    Envelope, ErrorCode, FrameError, Message, PROTOCOL_VERSION, SessionId, TurnDone, WireError,
    read_frame, write_frame,
};

/// Everything a turn needs, shared by every connection.
pub struct Context {
    pub cfg: Config,
    pub client: Client,
    pub port: Box<dyn ToolPort>,
    pub registry: Registry,
    busy: Mutex<HashSet<SessionId>>,
}

impl Context {
    pub fn new(
        cfg: Config,
        client: Client,
        port: Box<dyn ToolPort>,
        registry: Registry,
    ) -> Context {
        Context {
            cfg,
            client,
            port,
            registry,
            busy: Mutex::new(HashSet::new()),
        }
    }
}

/// While a turn runs, the session is in the busy set. The guard holds only an immutable borrow of
/// the context and a clone of the id, never the lock, so another connection can still check
/// membership while this turn is slow. The id is removed when the guard drops.
struct Held<'a> {
    ctx: &'a Context,
    id: SessionId,
}

impl Drop for Held<'_> {
    fn drop(&mut self) {
        let mut busy = self.ctx.busy.lock().unwrap_or_else(|p| p.into_inner());
        busy.remove(&self.id);
    }
}

/// Accepts connections forever, one thread each. Returns only if `accept` fails.
pub fn serve(listener: UnixListener, ctx: Arc<Context>) -> std::io::Result<()> {
    for stream in listener.incoming() {
        let stream = stream?;
        let ctx = Arc::clone(&ctx);
        std::thread::spawn(move || handle(stream, ctx));
    }
    Ok(())
}

/// One connection: read a turn, run it, stream its events, send the final frame.
pub fn handle(stream: UnixStream, ctx: Arc<Context>) {
    let mut stream = stream;

    // 1. Read the turn. A closed stream before anything is not an error; any other read failure
    //    is reported and closes the connection.
    let request = match read_frame(&mut stream) {
        Ok(env) => env,
        Err(FrameError::Closed) => return,
        Err(e) => {
            let code = match &e {
                FrameError::BadVersion(_) => ErrorCode::BadVersion,
                FrameError::Json(_) => ErrorCode::BadMessage,
                _ => ErrorCode::BadFrame,
            };
            let _ = write_frame(&mut stream, &error_frame(code, 0, e.to_string()));
            return;
        }
    };

    // 2. Only a turn is served; every reply after this carries the request's id and version.
    let Message::Turn(turn) = request.msg else {
        let _ = write_frame(
            &mut stream,
            &error_frame(ErrorCode::BadMessage, request.id, String::new()),
        );
        return;
    };

    // 3. Mark the session busy. If it already is, refuse at once. The guard releases the lock here
    //    and holds only the id plus an immutable borrow, so the turn below does not hold it.
    let held = {
        let mut busy = match ctx.busy.lock() {
            Ok(busy) => busy,
            Err(poisoned) => poisoned.into_inner(),
        };
        if busy.contains(&turn.session) {
            let _ = write_frame(
                &mut stream,
                &error_frame(ErrorCode::SessionBusy, request.id, String::new()),
            );
            return;
        }
        busy.insert(turn.session.clone());
        Held {
            ctx: &ctx,
            id: turn.session.clone(),
        }
    };

    // 4. Open or create the session, then run the turn, streaming each event as it happens.
    let runtime = Runtime {
        cfg: &ctx.cfg,
        client: &ctx.client,
        port: &*ctx.port,
        registry: &ctx.registry,
    };

    let mut session = if turn.resume {
        match Session::open(&ctx.cfg.paths.home, turn.session.clone()) {
            Ok(session) => session,
            Err(e) => {
                drop(held);
                let _ = write_frame(
                    &mut stream,
                    &error_frame(session_code(&e), request.id, e.to_string()),
                );
                return;
            }
        }
    } else {
        let baseline = match Baseline::assemble(&ctx.cfg, &ctx.registry) {
            Ok(baseline) => baseline,
            Err(e) => {
                drop(held);
                let _ = write_frame(
                    &mut stream,
                    &error_frame(ErrorCode::Internal, request.id, e.to_string()),
                );
                return;
            }
        };
        match Session::create(
            &ctx.cfg.paths.home,
            turn.session.clone(),
            baseline,
            ctx.cfg.slots.main,
        ) {
            Ok(session) => session,
            Err(e) => {
                drop(held);
                let _ = write_frame(
                    &mut stream,
                    &error_frame(session_code(&e), request.id, e.to_string()),
                );
                return;
            }
        }
    };

    let outcome = run_turn(&mut session, &runtime, &turn.content, &mut |event| {
        let _ = write_frame(
            &mut stream,
            &Envelope {
                v: PROTOCOL_VERSION,
                id: request.id,
                r#final: false,
                msg: Message::TurnEvent(event.clone()),
            },
        );
    });

    // 5. Release the session before the final frame, so the next turn does not find it busy.
    drop(held);

    // 6. and 7. A turn done is success; a turn error maps to a code. A write failure means the
    //    channel is gone, so we stop and finish quietly.
    match outcome {
        Ok(outcome) => {
            let _ = write_frame(
                &mut stream,
                &Envelope {
                    v: PROTOCOL_VERSION,
                    id: request.id,
                    r#final: true,
                    msg: Message::TurnDone(TurnDone {
                        content: outcome.content,
                        usage: outcome.usage,
                    }),
                },
            );
        }
        Err(e) => {
            let _ = write_frame(
                &mut stream,
                &error_frame(turn_code(&e), request.id, e.to_string()),
            );
        }
    }
}

/// One error envelope: final, the given id and version.
fn error_frame(code: ErrorCode, id: u64, detail: String) -> Envelope {
    Envelope {
        v: PROTOCOL_VERSION,
        id,
        r#final: true,
        msg: Message::Error(WireError { code, detail }),
    }
}

/// The channel code for a session that could not be opened or created.
fn session_code(e: &SessionError) -> ErrorCode {
    match e {
        SessionError::Exists(_) => ErrorCode::SessionExists,
        SessionError::NotFound(_) => ErrorCode::NoSuchSession,
        _ => ErrorCode::Internal,
    }
}

/// The channel code for a turn that ran but failed.
fn turn_code(e: &TurnError) -> ErrorCode {
    match e {
        TurnError::SessionFull => ErrorCode::SessionFull,
        TurnError::TurnLimit => ErrorCode::TurnLimit,
        TurnError::Infer(_) => ErrorCode::Inference,
        TurnError::Session(SessionError::Exists(_)) => ErrorCode::SessionExists,
        TurnError::Session(SessionError::NotFound(_)) => ErrorCode::NoSuchSession,
        TurnError::Session(_) => ErrorCode::Internal,
    }
}