use std::collections::{BTreeMap, BTreeSet};

use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::{
    runtime::RuntimeHandle,
    types::{
        DeliverySummaryRecord, TodoItem, WaitConditionSummary, WorkItemPlanArtifact,
        WorkItemPlanStatus, WorkItemReadiness, WorkItemRecord, WorkItemState,
    },
};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum WorkItemLifecycleView {
    Open,
    Completed,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum WorkItemFocusView {
    Current,
    Queued,
    Blocked,
    Completed,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub(crate) enum WorkItemCompletionReportSource {
    WorkItemResultSummary,
    DeliverySummary,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct WorkItemCompletionReportView {
    pub(crate) text: String,
    pub(crate) source: WorkItemCompletionReportSource,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) delivery_summary_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) source_turn_index: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) created_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct WorkItemView {
    pub(crate) id: String,
    pub(crate) agent_id: String,
    pub(crate) workspace_id: String,
    pub(crate) objective: String,
    pub(crate) state: WorkItemLifecycleView,
    pub(crate) focus: WorkItemFocusView,
    pub(crate) readiness: WorkItemReadiness,
    pub(crate) is_current: bool,
    pub(crate) is_runnable: bool,
    pub(crate) plan_status: WorkItemPlanStatus,
    pub(crate) plan_artifact: WorkItemPlanArtifact,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub(crate) todo_list: Vec<TodoItem>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) blocked_by: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) recheck_at: Option<DateTime<Utc>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) recheck_consumed_at: Option<DateTime<Utc>>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) completion_report: Option<WorkItemCompletionReportView>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub(crate) active_wait_conditions: Vec<WaitConditionSummary>,
    pub(crate) created_at: DateTime<Utc>,
    pub(crate) updated_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(crate) struct WorkItemQueryContext {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub(crate) current_work_item_id: Option<String>,
}

pub(crate) type WorkItemDeliverySummaryMap = BTreeMap<String, DeliverySummaryRecord>;
pub(crate) type WorkItemWaitConditionSummaryMap = BTreeMap<String, Vec<WaitConditionSummary>>;

pub(crate) async fn query_context(runtime: &RuntimeHandle) -> Result<WorkItemQueryContext> {
    let state = runtime.agent_state().await?;
    if let Some(bound_id) = state.current_turn_work_item_id.as_deref() {
        if let Some(record) = runtime.latest_work_item(bound_id).await? {
            if record.state == WorkItemState::Open {
                return Ok(WorkItemQueryContext {
                    current_work_item_id: Some(record.id),
                });
            }
        }
    }
    let current_work_item_id = match state.current_work_item_id.as_deref() {
        Some(current_id) => runtime
            .latest_work_item(current_id)
            .await?
            .filter(|record| record.state == WorkItemState::Open)
            .map(|record| record.id),
        None => None,
    };
    Ok(WorkItemQueryContext {
        current_work_item_id,
    })
}

pub(crate) async fn view_for_record(
    runtime: &RuntimeHandle,
    context: &WorkItemQueryContext,
    record: WorkItemRecord,
    include_todo_list: bool,
    delivery_summaries: Option<&WorkItemDeliverySummaryMap>,
    wait_conditions: Option<&WorkItemWaitConditionSummaryMap>,
) -> Result<WorkItemView> {
    let is_current = context.current_work_item_id.as_deref() == Some(record.id.as_str())
        && record.state == WorkItemState::Open;
    let mut record = record;
    crate::work_item_plan::refresh_plan_artifact_metadata(
        runtime.agent_home().as_path(),
        &mut record,
    )?;
    let plan_artifact = record
        .plan_artifact
        .clone()
        .ok_or_else(|| anyhow::anyhow!("missing plan artifact for work item {}", record.id))?;
    let todo_list = if include_todo_list {
        record.todo_list.clone()
    } else {
        Vec::new()
    };
    let state = lifecycle_view(&record.state);
    let focus = focus_view(&record, is_current);
    let readiness = record.readiness();
    let completion_report = completion_report_for_record(runtime, &record, delivery_summaries)?;
    let active_wait_conditions = match wait_conditions {
        Some(wait_conditions) => wait_conditions.get(&record.id).cloned().unwrap_or_default(),
        None => runtime
            .storage()
            .latest_active_wait_conditions_for_work_item(&record.agent_id, &record.id)?
            .into_iter()
            .map(WaitConditionSummary::from)
            .collect(),
    };
    Ok(WorkItemView {
        id: record.id,
        agent_id: record.agent_id,
        workspace_id: record.workspace_id,
        objective: record.objective,
        state,
        focus,
        readiness,
        is_current,
        is_runnable: readiness == WorkItemReadiness::Runnable,
        plan_status: record.plan_status,
        plan_artifact,
        todo_list,
        blocked_by: record.blocked_by,
        recheck_at: record.recheck_at,
        recheck_consumed_at: record.recheck_consumed_at,
        completion_report,
        active_wait_conditions,
        created_at: record.created_at,
        updated_at: record.updated_at,
    })
}

pub(crate) fn active_wait_conditions_by_work_item(
    runtime: &RuntimeHandle,
    records: &[WorkItemRecord],
) -> Result<WorkItemWaitConditionSummaryMap> {
    let agent_ids = records
        .iter()
        .map(|record| record.agent_id.as_str())
        .collect::<BTreeSet<_>>();
    let mut wait_conditions = BTreeMap::new();
    for agent_id in agent_ids {
        for condition in runtime
            .storage()
            .latest_active_wait_conditions_for_agent(agent_id)?
        {
            if let Some(work_item_id) = condition.work_item_id.as_ref() {
                wait_conditions
                    .entry(work_item_id.clone())
                    .or_insert_with(Vec::new)
                    .push(WaitConditionSummary::from(condition));
            }
        }
    }
    Ok(wait_conditions)
}

pub(crate) fn latest_delivery_summaries_by_work_item(
    runtime: &RuntimeHandle,
) -> Result<WorkItemDeliverySummaryMap> {
    let mut summaries = BTreeMap::new();
    for summary in runtime
        .storage()
        .read_recent_delivery_summaries(usize::MAX)?
        .into_iter()
        .rev()
        .filter(|summary| !summary.text.is_empty())
    {
        summaries
            .entry(summary.work_item_id.clone())
            .or_insert(summary);
    }
    Ok(summaries)
}

fn completion_report_for_record(
    runtime: &RuntimeHandle,
    record: &WorkItemRecord,
    delivery_summaries: Option<&WorkItemDeliverySummaryMap>,
) -> Result<Option<WorkItemCompletionReportView>> {
    if record.state != WorkItemState::Completed {
        return Ok(None);
    }
    let cached_delivery_summary = delivery_summaries
        .and_then(|summaries| summaries.get(&record.id))
        .cloned();
    let latest_delivery_summary = match (delivery_summaries.is_some(), cached_delivery_summary) {
        (_, Some(summary)) => Some(summary),
        (true, None) => None,
        (false, None) => runtime.storage().latest_delivery_summary(&record.id)?,
    };
    if let Some(text) = record
        .result_summary
        .as_ref()
        .filter(|text| !text.is_empty())
    {
        return Ok(Some(completion_report_view(
            text.clone(),
            WorkItemCompletionReportSource::WorkItemResultSummary,
            latest_delivery_summary
                .filter(|summary| summary.text == *text)
                .as_ref(),
        )));
    }
    Ok(latest_delivery_summary
        .filter(|summary| !summary.text.is_empty())
        .map(|summary| {
            completion_report_view(
                summary.text.clone(),
                WorkItemCompletionReportSource::DeliverySummary,
                Some(&summary),
            )
        }))
}

fn completion_report_view(
    text: String,
    source: WorkItemCompletionReportSource,
    delivery_summary: Option<&DeliverySummaryRecord>,
) -> WorkItemCompletionReportView {
    WorkItemCompletionReportView {
        text,
        source,
        delivery_summary_id: delivery_summary.map(|summary| summary.id.clone()),
        source_turn_index: delivery_summary.and_then(|summary| summary.source_turn_index),
        created_at: delivery_summary.map(|summary| summary.created_at),
    }
}

pub(crate) fn lifecycle_view(state: &WorkItemState) -> WorkItemLifecycleView {
    match state {
        WorkItemState::Open => WorkItemLifecycleView::Open,
        WorkItemState::Completed => WorkItemLifecycleView::Completed,
    }
}

pub(crate) fn focus_view(record: &WorkItemRecord, is_current: bool) -> WorkItemFocusView {
    if record.state == WorkItemState::Completed {
        return WorkItemFocusView::Completed;
    }
    if is_current {
        return WorkItemFocusView::Current;
    }
    if record.blocked_by.is_some() {
        WorkItemFocusView::Blocked
    } else {
        WorkItemFocusView::Queued
    }
}
