import json
import os
import re
from datetime import datetime
from tkinter import Tk, filedialog, messagebox, ttk, Label, Frame, StringVar, BooleanVar, colorchooser
from tkinterdnd2 import *
from pathlib import Path

def sanitize_filename(name):
    name = re.sub(r'[\\/*?:"<>|]', "_", name)
    name = re.sub(r'\s+', " ", name).strip()
    return name[:140]

def get_message_text(msg):
    if not msg: return ""
    if isinstance(msg.get("text"), str): return msg["text"]
    if isinstance(msg.get("content"), list):
        texts = [item.get("text", str(item)) if isinstance(item, dict) else str(item) for item in msg["content"]]
        return "\n\n".join(filter(None, texts))
    return str(msg.get("content", ""))

def extract_project_name(conv):
    project = conv.get("project") or conv.get("workspace") or conv.get("project_name")
    if project and isinstance(project, dict):
        project = project.get("name") or project.get("title")
    return project if project else None

def convert_urls_to_links(text):
    url_pattern = r'(https?://[^\s<>"\']+)'
    return re.sub(url_pattern, r'<a href="\1" target="_blank">\1</a>', text)

class ClaudeExporterGUI:
    def __init__(self):
        self.root = TkinterDnD.Tk()
        self.root.title("Claude Exported Data Offline Extractor & Browser")
        self.root.geometry("852x612")      # New initial size
        self.root.minsize(680, 580)
        self.root.resizable(True, True)
        
        script_dir = os.path.dirname(os.path.abspath(__file__))
        self.input_file = StringVar(value=os.path.join(script_dir, "conversations.json"))
        self.output_dir = StringVar(value=script_dir)
        
        self.export_md = BooleanVar(value=True)
        self.export_html = BooleanVar(value=True)
        self.user_color = StringVar(value="#E3F2FD")
        self.claude_color = StringVar(value="#F1F8E9")
        
        self.create_widgets()
        
    def create_widgets(self):
        main_frame = Frame(self.root, padx=25, pady=20)
        main_frame.pack(fill="both", expand=True)
        
        # Title
        title = Label(main_frame, text="Claude Exported Data Offline Extractor & Browser", 
                      font=("Arial", 16, "bold"))
        title.pack(pady=(0, 10))
        
        # Description
        desc_text = "This is a tool for offline unpacking and browsing exported Claude conversations stored in conversations.json " \
                    "file which can be found in the zip file sent to your email after selecting " \
                    "\"Settings → Privacy → Export data\" in your Claude"
        
        self.desc_label = Label(main_frame, text=desc_text, font=("Arial", 11), 
                               justify="center", fg="black", wraplength=720)
        self.desc_label.pack(pady=(0, 20))
        
        # Input File
        Label(main_frame, text="Input File (conversations.json) :").pack(anchor="w")
        file_frame = Frame(main_frame)
        file_frame.pack(fill="x", pady=4)
        self.input_entry = ttk.Entry(file_frame, textvariable=self.input_file)
        self.input_entry.pack(side="left", fill="x", expand=True)
        ttk.Button(file_frame, text="Browse", command=self.browse_input).pack(side="right", padx=5)
        
        self.input_entry.drop_target_register(DND_FILES)
        self.input_entry.dnd_bind('<<Drop>>', self.drop_file)
        
        # Output Folder
        Label(main_frame, text="Output Folder :").pack(anchor="w", pady=(10,0))
        out_frame = Frame(main_frame)
        out_frame.pack(fill="x", pady=4)
        ttk.Entry(out_frame, textvariable=self.output_dir).pack(side="left", fill="x", expand=True)
        ttk.Button(out_frame, text="Browse", command=self.browse_output).pack(side="right", padx=5)
        
        # Export Options
        opt_frame = Frame(main_frame)
        opt_frame.pack(fill="x", pady=12)
        ttk.Checkbutton(opt_frame, text="Export as Markdown (.md)", variable=self.export_md).pack(anchor="w")
        ttk.Checkbutton(opt_frame, text="Export as HTML (with colors, navigation & clickable links)", 
                       variable=self.export_html).pack(anchor="w", pady=4)
        
        # Colors (Centered)
        color_frame = Frame(main_frame)
        color_frame.pack(pady=15)
        
        Label(color_frame, text="User background:").pack(side="left", padx=8)
        ttk.Button(color_frame, text="Pick", command=self.pick_user_color).pack(side="left", padx=5)
        Label(color_frame, textvariable=self.user_color, width=10, background="#E3F2FD").pack(side="left", padx=8)
        
        Label(color_frame, text="Claude background:").pack(side="left", padx=15)
        ttk.Button(color_frame, text="Pick", command=self.pick_claude_color).pack(side="left", padx=5)
        Label(color_frame, textvariable=self.claude_color, width=10, background="#F1F8E9").pack(side="left", padx=8)
        
        # Start Button
        ttk.Button(main_frame, text="Start Export", command=self.start_conversion).pack(pady=20, fill="x")
        
        self.progress = ttk.Progressbar(main_frame, mode='determinate')
        self.progress.pack(fill="x", pady=5)
        
        self.status = Label(main_frame, text="", foreground="gray")
        self.status.pack(pady=5)
        
        # Make description wrap dynamically
        self.root.bind("<Configure>", self.on_resize)
        
    def on_resize(self, event=None):
        if hasattr(self, 'desc_label'):
            width = self.root.winfo_width() - 80
            self.desc_label.config(wraplength=max(400, width))
    
    def drop_file(self, event):
        path = event.data
        if path.startswith('{') and path.endswith('}'):
            path = path[1:-1]
        if path.lower().endswith('.json'):
            self.input_file.set(path)
            self.status.config(text=f"Loaded: {Path(path).name}")
    
    def pick_user_color(self):
        color = colorchooser.askcolor()[1]
        if color: self.user_color.set(color)
    
    def pick_claude_color(self):
        color = colorchooser.askcolor()[1]
        if color: self.claude_color.set(color)
    
    def browse_input(self):
        file = filedialog.askopenfilename(filetypes=[("JSON files", "*.json")])
        if file: self.input_file.set(file)
    
    def browse_output(self):
        folder = filedialog.askdirectory()
        if folder: self.output_dir.set(folder)
    
    def start_conversion(self):
        if not os.path.exists(self.input_file.get()):
            messagebox.showerror("Error", "conversations.json file not found!")
            return
        
        self.progress['value'] = 0
        self.status.config(text="Processing...", foreground="blue")
        self.root.update()
        
        try:
            success = self.convert()
            messagebox.showinfo("Success", f"Successfully exported {success} conversations!")
            self.status.config(text="Export completed successfully ✓", foreground="green")
        except Exception as e:
            messagebox.showerror("Error", f"An error occurred:\n{str(e)}")
            self.status.config(text="Error during export", foreground="red")
        finally:
            self.progress['value'] = 0
    
    # === Conversion methods (same as before) ===
    def convert(self):
        with open(self.input_file.get(), "r", encoding="utf-8") as f:
            data = json.load(f)
        
        convs = data.get("conversations", data) if isinstance(data, dict) else data
        if not isinstance(convs, list):
            convs = [convs]
        
        out_dir = self.output_dir.get()
        os.makedirs(out_dir, exist_ok=True)
        
        success = 0
        self.progress['maximum'] = len(convs)
        months = {}
        
        for idx, conv in enumerate(convs):
            conv_number = f"{idx+1:03d}"
            title = conv.get("title") or f"Conversation {conv_number}"
            project_name = extract_project_name(conv)
            created_at = conv.get("create_time") or conv.get("created_at")
            
            date_str = "unknown_date"
            month_key = "unknown"
            if created_at:
                try:
                    dt = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
                    date_str = dt.strftime("%Y-%m-%d")
                    month_key = dt.strftime("%Y-%m")
                except:
                    pass
            
            display_title = f"Claude Exported Conversation {conv_number}"
            if project_name:
                display_title += f" - {project_name}"
            
            if self.export_md.get():
                self.save_markdown(conv, display_title, date_str, out_dir, project_name)
            
            if self.export_html.get():
                html_rel_path = self.save_html(conv, display_title, date_str, out_dir, month_key, 
                                            self.user_color.get(), self.claude_color.get())
                if month_key not in months:
                    months[month_key] = []
                months[month_key].append((display_title, date_str, html_rel_path))
            
            success += 1
            self.progress['value'] = idx + 1
            self.root.update_idletasks()
        
        if self.export_html.get():
            self.create_main_html_index(out_dir, months)
        
        return success
    
    def save_markdown(self, conv, title, date_str, base_dir, project_name):
        md = [f"# {title}\n", f"**Created:** {date_str}"]
        if project_name:
            md.append(f"**Project:** {project_name}")
        md.append("\n---\n\n")
        
        messages = conv.get("chat_messages") or conv.get("messages") or []
        for msg in messages:
            sender = msg.get("sender") or msg.get("role") or "unknown"
            role = "You" if sender in ["human", "user"] else "Claude" if sender in ["assistant", "claude"] else sender.capitalize()
            text = get_message_text(msg)
            if text.strip():
                md.append(f"### {role}\n")
                md.append(f"{text}\n")
                md.append("---\n")
        
        safe_title = sanitize_filename(title)
        filepath = os.path.join(base_dir, f"{date_str}_{safe_title}.md")
        with open(filepath, "w", encoding="utf-8") as f:
            f.write("\n".join(md))
    
    def save_html(self, conv, title, date_str, base_dir, month_key, user_color, claude_color):
        html = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>{title}</title>
    <style>
        body {{font-family: Consolas, monospace; max-width: 960px; margin: 40px auto; padding: 20px; line-height: 1.65; text-align: justify;}}
        h1 {{color: #1a73e8; text-align: center;}}
        .back-link {{text-align: left; margin-bottom: 25px;}}
        .message {{margin: 28px 0; padding: 18px; border-radius: 10px;}}
        .user {{background: {user_color};}}
        .claude {{background: {claude_color};}}
        .speaker {{font-weight: bold; margin-bottom: 12px; font-size: 1.1em;}}
        a {{color: #1a73e8;}}
    </style>
</head>
<body>
    <div class="back-link"><a href="INDEX.html">← Back to Index</a></div>
    <h1>{title}</h1>
    <p><strong>Created:</strong> {date_str}</p>
    <hr>"""
        
        messages = conv.get("chat_messages") or conv.get("messages") or []
        for msg in messages:
            sender = msg.get("sender") or msg.get("role") or "unknown"
            role = "You" if sender in ["human", "user"] else "Claude"
            raw_text = get_message_text(msg)
            text = convert_urls_to_links(raw_text.replace("\n", "<br>"))
            css_class = "user" if role == "You" else "claude"
            html += f'<div class="message {css_class}"><div class="speaker">{role}</div>{text}</div>'
        
        html += "</body></html>"
        
        safe_title = sanitize_filename(title)
        filepath = os.path.join(base_dir, f"{date_str}_{safe_title}.html")
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(html)
        return f"{date_str}_{safe_title}.html"
    
    def create_main_html_index(self, base_dir, months):
        index_html = """<!DOCTYPE html>
<html><head><meta charset="utf-8"><title>Claude Archive Index</title>
<style>body{font-family:Consolas,monospace;margin:40px;line-height:1.6;} h1{color:#1a73e8;} h2{margin-top:35px;} ul{list-style:none;} li{margin:10px 0;}</style>
</head><body><h1>Claude Exported Conversations Archive</h1>"""
        
        for month in sorted(months.keys(), reverse=True):
            index_html += f"<h2>{month}</h2><ul>"
            for title, date, link in months[month]:
                index_html += f'<li><a href="{link}">{date} — {title}</a></li>'
            index_html += "</ul>"
        
        index_html += "</body></html>"
        
        with open(os.path.join(base_dir, "INDEX.html"), "w", encoding="utf-8") as f:
            f.write(index_html)

if __name__ == "__main__":
    app = ClaudeExporterGUI()
    app.root.mainloop()