use crate::proto::cortex_client::CortexClient;
use crate::proto::*;
use rice_core::error::RiceError;
use std::str::FromStr;
use tonic::codegen::InterceptedService;
use tonic::metadata::MetadataValue;
use tonic::service::Interceptor;
use tonic::transport::Channel;
use tonic::Status;

#[derive(Clone)]
struct AuthInterceptor {
    token: MetadataValue<tonic::metadata::Ascii>,
}

impl Interceptor for AuthInterceptor {
    fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, Status> {
        req.metadata_mut()
            .insert("authorization", self.token.clone());
        Ok(req)
    }
}

#[derive(Clone)]
enum StateClientInner {
    Unauth(CortexClient<Channel>),
    Auth(CortexClient<InterceptedService<Channel, AuthInterceptor>>),
}

#[derive(Clone)]
pub struct StateClient {
    inner: StateClientInner,
}

impl StateClient {
    pub async fn connect(dst: String) -> Result<Self, RiceError> {
        let client = CortexClient::connect(dst).await?;
        Ok(Self {
            inner: StateClientInner::Unauth(client),
        })
    }

    pub async fn connect_with_token(dst: String, token: String) -> Result<Self, RiceError> {
        let channel = Channel::from_shared(dst)
            .map_err(|e| RiceError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e)))?
            .connect()
            .await?;

        let token_meta = MetadataValue::from_str(&format!("Bearer {}", token))
            .map_err(|e| RiceError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e)))?;

        let interceptor = AuthInterceptor { token: token_meta };
        let client = CortexClient::with_interceptor(channel, interceptor);

        Ok(Self {
            inner: StateClientInner::Auth(client),
        })
    }

    pub async fn focus(
        &mut self,
        content: String,
        run_id: String,
    ) -> Result<FocusResponse, RiceError> {
        let request = FocusRequest { content, run_id };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.focus(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.focus(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn drift(&mut self, run_id: String) -> Result<DriftResponse, RiceError> {
        let request = DriftRequest { run_id };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.drift(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.drift(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn commit(&mut self, trace: Trace) -> Result<Ack, RiceError> {
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.commit(trace).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.commit(trace).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn set_variable(
        &mut self,
        run_id: String,
        name: String,
        value_json: String,
        source: String,
    ) -> Result<Ack, RiceError> {
        let request = SetVariableRequest {
            run_id,
            name,
            value_json,
            source,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.set_variable(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.set_variable(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn get_variable(
        &mut self,
        run_id: String,
        name: String,
    ) -> Result<VariableResponse, RiceError> {
        let request = GetVariableRequest { run_id, name };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.get_variable(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.get_variable(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn list_variables(
        &mut self,
        run_id: String,
    ) -> Result<ListVariablesResponse, RiceError> {
        let request = ListVariablesRequest { run_id };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.list_variables(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.list_variables(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn delete_variable(
        &mut self,
        run_id: String,
        name: String,
    ) -> Result<Ack, RiceError> {
        let request = DeleteVariableRequest { run_id, name };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.delete_variable(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.delete_variable(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn add_goal(
        &mut self,
        run_id: String,
        description: String,
        priority: String,
        parent_id: Option<String>,
    ) -> Result<GoalResponse, RiceError> {
        let request = AddGoalRequest {
            run_id,
            description,
            priority,
            parent_id: parent_id.unwrap_or_default(),
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.add_goal(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.add_goal(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn update_goal(
        &mut self,
        run_id: String,
        goal_id: String,
        status: String,
    ) -> Result<Ack, RiceError> {
        let request = UpdateGoalRequest {
            run_id,
            goal_id,
            status,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.update_goal(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.update_goal(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn list_goals(
        &mut self,
        run_id: String,
        status_filter: Option<String>,
    ) -> Result<ListGoalsResponse, RiceError> {
        let request = ListGoalsRequest {
            run_id,
            status_filter: status_filter.unwrap_or_default(),
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.list_goals(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.list_goals(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn submit_action(
        &mut self,
        run_id: String,
        agent_id: String,
        action_type: String,
        action_json: String,
    ) -> Result<ActionResponse, RiceError> {
        let request = ActionRequest {
            run_id,
            agent_id,
            action_type,
            action_json,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.submit_action(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.submit_action(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn get_action_log(
        &mut self,
        run_id: String,
        limit: u64,
        action_type_filter: Option<String>,
    ) -> Result<ActionLogResponse, RiceError> {
        let request = ActionLogRequest {
            run_id,
            limit,
            action_type_filter: action_type_filter.unwrap_or_default(),
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.get_action_log(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.get_action_log(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn reminisce(
        &mut self,
        embedding: Vec<f32>,
        limit: u64,
        query_text: String,
        run_id: String,
    ) -> Result<RecallResponse, RiceError> {
        let request = RecallRequest {
            embedding,
            limit,
            query_text,
            filter: "".to_string(),
            run_id,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.reminisce(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.reminisce(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn define_concept(
        &mut self,
        run_id: String,
        name: String,
        schema_json: String,
    ) -> Result<Ack, RiceError> {
        let request = DefineConceptRequest {
            run_id,
            name,
            schema_json,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.define_concept(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.define_concept(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn list_concepts(
        &mut self,
        run_id: String,
    ) -> Result<ListConceptsResponse, RiceError> {
        let request = ListConceptsRequest { run_id };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.list_concepts(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.list_concepts(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn run_cycle(
        &mut self,
        run_id: String,
        agent_id: String,
        candidates: Vec<ActionCandidate>,
    ) -> Result<CycleResponse, RiceError> {
        let request = RunCycleRequest {
            run_id,
            agent_id,
            candidates,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.run_cycle(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.run_cycle(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn get_cycle_history(
        &mut self,
        run_id: String,
        limit: u64,
    ) -> Result<CycleHistoryResponse, RiceError> {
        let request = CycleHistoryRequest { run_id, limit };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.get_cycle_history(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.get_cycle_history(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn delete_run(&mut self, run_id: String) -> Result<Ack, RiceError> {
        let request = RunRequest { run_id };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.delete_run(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.delete_run(request).await?;
                Ok(response.into_inner())
            }
        }
    }

    pub async fn subscribe(
        &mut self,
        run_id: String,
        event_types: Vec<String>,
    ) -> Result<tonic::Streaming<Event>, RiceError> {
        let request = SubscribeRequest {
            run_id,
            event_types,
        };
        match &mut self.inner {
            StateClientInner::Unauth(client) => {
                let response = client.subscribe(request).await?;
                Ok(response.into_inner())
            }
            StateClientInner::Auth(client) => {
                let response = client.subscribe(request).await?;
                Ok(response.into_inner())
            }
        }
    }
}
