gway 0.1.3

use std::{net::SocketAddr, sync::Arc};

use http::{Request, Response, StatusCode, header, response::Builder};
use http_body::Body;
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use hyper::{body::Bytes, client::conn::http1};
use hyper_util::rt::TokioIo;
use sub_host::sub_host;
use tokio::net::TcpStream;

use crate::{Error, Result, Route, req_host};

// 1. Define the new trait
pub trait IntoError {
  fn into_error(self) -> Error;
}

// 2. Implement it for specific error types
impl IntoError for hyper::Error {
  fn into_error(self) -> Error {
    Error::Hyper(self)
  }
}

impl IntoError for std::convert::Infallible {
  fn into_error(self) -> Error {
    Error::Infallible(self)
  }
}

pub static mut N: usize = 0;

pub async fn fetch(
  upstream_addr: SocketAddr,
  req: Request<Full<Bytes>>,
) -> Result<BoxBody<Bytes, hyper::Error>> {
  let stream = TcpStream::connect(upstream_addr).await?;
  let io = TokioIo::new(stream);
  let (mut sender, conn) = http1::handshake(io).await?;

  tokio::task::spawn(async move {
    if let Err(err) = conn.await {
      eprintln!("{upstream_addr} {err}");
    }
  });

  Ok(sender.send_request(req).await.map(|b| b.boxed())?)
}

pub async fn proxy<B>(req: Request<B>, route: Arc<Route>) -> Response<BoxBody<Bytes, hyper::Error>>
where
  B: Body<Data = Bytes> + Send + 'static,
  B::Error: IntoError + Send + Sync + 'static,
{
  let host = req_host(&req).to_owned();
  let path = req
    .uri()
    .path_and_query()
    .map(|x| x.as_str())
    .unwrap_or("")
    .to_owned();
  match _proxy(&host, &path, req, route).await {
    Ok(res) => {
      let status = res.status();
      println!("{status} {host} {path}");
      res
    }
    Err(err) => {
      let err = err.to_string();
      eprintln!("Error: {host} {path} {err}");
      response(|b| b.status(500), err).unwrap_or_default()
    }
  }
}

fn response(
  build: impl Fn(Builder) -> Builder,
  body: impl Into<Bytes>,
) -> Result<Response<BoxBody<Bytes, hyper::Error>>> {
  Ok(
    build(Builder::new()).body(
      Full::new(body.into())
        .map_err(|never| match never {})
        .boxed(),
    )?,
  )
}

// 3. Update the function signature
pub async fn _proxy<B>(
  host: &str,
  path_and_query: &str,
  req: Request<B>,
  route: Arc<Route>,
) -> Result<Response<BoxBody<Bytes, hyper::Error>>>
where
  B: Body<Data = Bytes> + Send + 'static,
  B::Error: IntoError + Send + Sync + 'static,
{
  if let Some(site_conf) = route.host_conf.get(host) {
    let site_conf = site_conf.value();
    let upstream = &site_conf.upstream;
    let upstream_addr_li = &upstream.addr_li;
    let len = upstream_addr_li.len();
    if len == 0 {
      return Err(Error::UpstreamNotFound);
    }

    let (parts, body) = req.into_parts();
    // 4. Update the error handling
    let body = body.collect().await.map_err(|e| e.into_error())?.to_bytes();

    let mut retry = 0;
    let mut pos = unsafe {
      N = N.overflowing_add(1).0;
      N
    } % len;
    loop {
      let upstream_addr = upstream_addr_li[pos];
      let req = Request::from_parts(parts.clone(), Full::new(body.clone()));
      match fetch(upstream_addr, req).await {
        Ok(res) => {
          return Ok(Response::new(res));
        }
        Err(err) => {
          eprintln!("Error: {host} {path_and_query} {upstream_addr} {}", err);
          retry += 1;
          if retry > upstream.max_retry {
            return Err(err);
          }
          pos = (pos + 1) % len;
        }
      }
    }
  } else {
    if let Some(host) = sub_host(host)
      && route.host_conf.get(&faststr::FastStr::new(&host)).is_some()
    {
      return response(
        |b| {
          let new_uri = format!("https://{}{}", host, path_and_query);
          b.status(StatusCode::MOVED_PERMANENTLY)
            .header(header::LOCATION, new_uri)
        },
        &b""[..],
      );
    }
    response(|b| b.status(404), &b"404: Not Found"[..])
  }
}