import json
import re
from datetime import datetime
from pathlib import Path

from docx import Document
from docx.shared import Pt, Inches, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH, WD_BREAK
from docx.enum.style import WD_STYLE_TYPE

class DeepSeekExporter:
    def __init__(self, json_file_path):
        with open(json_file_path, 'r', encoding='utf-8') as f:
            self.data = json.load(f)
        self.doc = Document()
        self._setup_styles()

    def _setup_styles(self):
        """Setup custom styles for better formatting"""
        # Title style
        title_style = self.doc.styles['Title']
        title_style.font.size = Pt(24)
        title_style.font.bold = True
        title_style.font.color.rgb = RGBColor(0, 102, 204)

        # Heading 1 for conversation titles
        h1 = self.doc.styles['Heading 1']
        h1.font.size = Pt(18)
        h1.font.bold = True
        h1.font.color.rgb = RGBColor(51, 51, 51)

        # User message style
        user_style = self.doc.styles.add_style('UserMessage', WD_STYLE_TYPE.PARAGRAPH)
        user_style.font.name = 'Calibri'
        user_style.font.size = Pt(11)
        user_style.font.color.rgb = RGBColor(0, 0, 139)  # Dark blue
        user_style.paragraph_format.left_indent = Inches(0.3)
        user_style.paragraph_format.space_after = Pt(6)

        # Assistant message style
        assistant_style = self.doc.styles.add_style('AssistantMessage', WD_STYLE_TYPE.PARAGRAPH)
        assistant_style.font.name = 'Calibri'
        assistant_style.font.size = Pt(11)
        assistant_style.font.color.rgb = RGBColor(0, 0, 0)
        assistant_style.paragraph_format.left_indent = Inches(0.3)
        assistant_style.paragraph_format.space_after = Pt(6)

        # Thinking/Reasoning style (gray, italic)
        thinking_style = self.doc.styles.add_style('ThinkingBlock', WD_STYLE_TYPE.PARAGRAPH)
        thinking_style.font.name = 'Calibri'
        thinking_style.font.size = Pt(10)
        thinking_style.font.italic = True
        thinking_style.font.color.rgb = RGBColor(80, 80, 80)
        thinking_style.paragraph_format.left_indent = Inches(0.5)
        thinking_style.paragraph_format.space_after = Pt(4)

        # Code block style
        code_style = self.doc.styles.add_style('CodeBlock', WD_STYLE_TYPE.PARAGRAPH)
        code_style.font.name = 'Consolas'
        code_style.font.size = Pt(9)
        code_style.font.color.rgb = RGBColor(0, 0, 0)
        code_style.paragraph_format.left_indent = Inches(0.5)
        code_style.paragraph_format.space_after = Pt(4)

    def _clean_text(self, text):
        """Clean up text formatting"""
        if not text:
            return ""
        # Remove excessive whitespace
        text = re.sub(r'\n{3,}', '\n\n', text)
        return text.strip()

    def _add_fragment(self, fragment):
        """Add a single message fragment to the document"""
        frag_type = fragment.get('type')
        content = fragment.get('content', '')

        if not content:
            return

        if frag_type == 'REQUEST':
            p = self.doc.add_paragraph(style='UserMessage')
            p.paragraph_format.left_indent = Inches(0)
            # Add a nice "User" label
            run = p.add_run('👤 ')
            run.bold = True
            run.font.size = Pt(11)
            run.font.color.rgb = RGBColor(0, 0, 139)
            run = p.add_run(content)
            run.font.size = Pt(11)

        elif frag_type == 'THINK':
            p = self.doc.add_paragraph(style='ThinkingBlock')
            p.paragraph_format.left_indent = Inches(0.5)
            run = p.add_run('🧠 Thinking...\n')
            run.bold = True
            run.font.size = Pt(9)
            run.font.color.rgb = RGBColor(80, 80, 80)
            run = p.add_run(content[:1000] + ('...' if len(content) > 1000 else ''))
            run.font.size = Pt(10)
            run.font.color.rgb = RGBColor(80, 80, 80)

        elif frag_type == 'RESPONSE':
            p = self.doc.add_paragraph(style='AssistantMessage')
            p.paragraph_format.left_indent = Inches(0)
            # Add a nice "Assistant" label
            run = p.add_run('🤖 ')
            run.bold = True
            run.font.size = Pt(11)
            run.font.color.rgb = RGBColor(0, 102, 0)
            # Clean and add content
            content = self._clean_text(content)
            for line in content.split('\n'):
                if line.strip().startswith('```'):
                    # Handle code blocks
                    self.doc.add_paragraph(line.strip(), style='CodeBlock')
                elif line.strip().startswith('**'):
                    # Bold headers
                    p = self.doc.add_paragraph()
                    run = p.add_run(line.strip())
                    run.bold = True
                    run.font.size = Pt(12)
                elif re.match(r'^[\d]+\.', line.strip()):
                    # Numbered lists
                    p = self.doc.add_paragraph(line.strip(), style='List Number')
                elif line.strip().startswith('- ') or line.strip().startswith('* '):
                    # Bullet points
                    p = self.doc.add_paragraph(line.strip(), style='List Bullet')
                elif line.strip() and not line.strip().startswith('---'):
                    run = p.add_run(line.strip())
                    run.font.size = Pt(11)

        elif frag_type == 'SEARCH':
            # Handle search results
            results = fragment.get('results', [])
            if results:
                p = self.doc.add_paragraph()
                p.paragraph_format.left_indent = Inches(0.3)
                run = p.add_run('🔍 Search results:\n')
                run.bold = True
                run.font.size = Pt(9)
                run.font.color.rgb = RGBColor(128, 128, 128)
                for result in results[:3]:
                    url = result.get('url', '')
                    title = result.get('title', '')
                    if title:
                        p = self.doc.add_paragraph(f'- {title}: {url}', style='CodeBlock')

    def _process_mapping(self, mapping):
        """Process a conversation's mapping structure to get ordered messages"""
        # Build message order by following parent-child links
        messages = []
        current_id = 'root'

        # Get all message nodes
        message_nodes = {}
        for node_id, node_data in mapping.items():
            if node_id != 'root' and node_data.get('message'):
                message_nodes[node_id] = node_data

        # Follow the chain from root's first child
        root_children = mapping.get('root', {}).get('children', [])
        if not root_children:
            return messages

        current_id = root_children[0]
        while current_id:
            node = mapping.get(current_id)
            if not node:
                break
            if node.get('message'):
                messages.append(node)
            children = node.get('children', [])
            if children:
                current_id = children[0]
            else:
                current_id = None

        return messages

    def export(self, output_path="deepseek_export.docx"):
        """Main export function"""
        # Add title
        title = self.doc.add_heading('DeepSeek Conversations', level=0)
        title.alignment = WD_ALIGN_PARAGRAPH.CENTER
        self.doc.add_paragraph(f'Generated: {datetime.now().strftime("%B %d, %Y at %I:%M %p")}')
        self.doc.add_paragraph()
        self.doc.add_page_break()

        # Process each conversation
        for idx, conversation in enumerate(self.data, 1):
            # Conversation header
            title_text = conversation.get('title', f'Conversation {idx}')
            self.doc.add_heading(title_text, level=1)

            # Metadata
            inserted = conversation.get('inserted_at', '')
            updated = conversation.get('updated_at', '')
            if inserted:
                try:
                    dt = datetime.fromisoformat(inserted.replace('+08:00', ''))
                    self.doc.add_paragraph(f'📅 Started: {dt.strftime("%B %d, %Y at %I:%M %p")}')
                except:
                    pass

            self.doc.add_paragraph('─' * 60)
            self.doc.add_paragraph()

            # Process messages in order
            mapping = conversation.get('mapping', {})
            messages = self._process_mapping(mapping)

            for msg_node in messages:
                msg_data = msg_node.get('message', {})
                fragments = msg_data.get('fragments', [])

                # Sort fragments by type (REQUEST first, then THINK, SEARCH, RESPONSE)
                order = {'REQUEST': 0, 'SEARCH': 1, 'THINK': 2, 'RESPONSE': 3}
                fragments.sort(key=lambda x: order.get(x.get('type'), 99))

                for fragment in fragments:
                    self._add_fragment(fragment)

                # Add a separator between messages
                self.doc.add_paragraph()
                self.doc.add_paragraph('—' * 40, style='No Spacing')

            # Add page break between conversations (except last one)
            if idx < len(self.data):
                self.doc.add_page_break()

        # Save the document
        self.doc.save(output_path)
        print(f"✅ Document saved as: {output_path}")
        print(f"📊 Exported {len(self.data)} conversations")

        return output_path


if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python deepseek_to_word.py <exported_chats.json> [output.docx]")
        print("Example: python deepseek_to_word.py deepseek_export.json my_conversations.docx")
        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else "deepseek_export.docx"

    if not Path(input_file).exists():
        print(f"❌ Error: File '{input_file}' not found!")
        sys.exit(1)

    exporter = DeepSeekExporter(input_file)
    exporter.export(output_file)