.assistant-message h1, .assistant-message h2, .assistant-message h3 {
  margin-top: 1.5em;
  margin-bottom: 0.5em;
  color: var(--text-color);
}

.assistant-message p {
  margin-bottom: 1em;
}

.assistant-message code:not(pre code) {
  background-color: rgba(255, 255, 255, 0.1);
  padding: 0.2em 0.4em;
  border-radius: 4px;
  font-family: 'Courier New', Courier, monospace;
  font-size: 0.9em;
}

.code-block-wrapper {
  margin: 1rem 0;
  border-radius: 0.5rem;
  overflow: hidden;
  border: 1px solid var(--border-color);
  background-color: var(--code-bg);
}

.code-header {
  display: flex;
  justify-content: space-between;
  align-items: center;
  padding: 0.5rem 1rem;
  background-color: #1f1f1f;
  font-size: 0.8rem;
  color: var(--muted-text-color);
  border-bottom: 1px solid var(--border-color);
}

.code-header button {
  background: transparent;
  border: none;
  color: inherit;
  cursor: pointer;
  display: flex;
  align-items: center;
  gap: 0.25rem;
}

.code-header button:hover {
  color: var(--text-color);
}

pre {
  margin: 0;
  padding: 1rem;
  overflow-x: auto;
}

pre code {
  font-family: 'Courier New', Courier, monospace;
  background: transparent !important;
}

pre::-webkit-scrollbar {
  height: 8px;
}

pre::-webkit-scrollbar-track {
  background: #0a0a0a;
  border-radius: 4px;
}

pre::-webkit-scrollbar-thumb {
  background: #404040;
  border-radius: 4px;
}

pre::-webkit-scrollbar-thumb:hover {
  background: #505050;
}

.code-collapsed pre {
  max-height: 150px;
  overflow-y: hidden;
  mask-image: linear-gradient(to bottom, black 50%, transparent 100%);
  -webkit-mask-image: linear-gradient(to bottom, black 50%, transparent 100%);
}

.code-collapsed.expanded pre {
  max-height: none;
  overflow-y: visible;
  mask-image: none;
  -webkit-mask-image: none;
}

.expand-code-btn {
  position: relative;
  background: rgba(255, 255, 255, 0.05);
  border: 1px solid var(--border-color);
  padding: 0.5rem 1rem;
  border-radius: 0.4rem;
  color: var(--muted-text-color);
  cursor: pointer;
  font-size: 0.75rem;
  width: 100%;
  margin-top: -1px;
  border-top: none;
  transition: all 0.2s;
  display: flex;
  align-items: center;
  justify-content: center;
  gap: 0.5rem;
}

.expand-code-btn:hover {
  background: rgba(255, 255, 255, 0.1);
  color: var(--text-color);
}

.copy-button {
  float: right;
  margin-left: 0.5rem;
  cursor: pointer;
  background: transparent;
  border: 1px solid var(--border-color);
  color: var(--muted-text-color);
  border-radius: 4px;
  padding: 2px 6px;
  font-size: 0.75rem;
  transition: all 0.2s;
}

.expand-btn {
  width: 100%;
  padding: 0.5rem;
  background-color: #1f1f1f;
  border: none;
  border-top: 1px solid var(--border-color);
  color: var(--muted-text-color);
  cursor: pointer;
  font-size: 0.8rem;
}

.expand-btn:hover {
  background-color: var(--hover-color);
}