import streamlit as st
import re


def edit_model_form(model_info):
    """Form for editing model metadata"""
    st.subheader("Edit Model Information")

    if not model_info:
        st.error("Model information not found")
        return False, None

    # Extract model card content (README.md) if available
    model_card_content = ""
    try:
        repo_id = model_info.modelId
        model_card_url = f"https://huggingface.co/{repo_id}/raw/main/README.md"
        response = st.session_state.client.api._get_paginated(model_card_url)
        if response.status_code == 200:
            model_card_content = response.text
    except Exception as e:
        st.warning(f"Couldn't load model card: {str(e)}")

    # Extract tags from model card
    tags = []
    if model_card_content:
        # Look for tags section in YAML frontmatter
        yaml_match = re.search(r"---\s+(.*?)\s+---", model_card_content, re.DOTALL)
        if yaml_match:
            yaml_content = yaml_match.group(1)
            tags_match = re.search(r"tags:\s*((?:- .*?\n)+)", yaml_content, re.DOTALL)
            if tags_match:
                tags_content = tags_match.group(1)
                tags = [
                    line.strip("- \n")
                    for line in tags_content.split("\n")
                    if line.strip().startswith("-")
                ]

    # Extract description (first paragraph after the title)
    description = ""
    if model_card_content:
        # Find content after title and before next heading
        title_match = re.search(
            r"# .*?\n\n(.*?)(?=\n## |\Z)", model_card_content, re.DOTALL
        )
        if title_match:
            description = title_match.group(1).strip()

    with st.form("edit_model_form"):
        # Model tags
        st.markdown("#### Model Tags")
        available_tags = st.session_state.client.get_model_tags()
        selected_tags = st.multiselect(
            "Select tags for your model",
            options=available_tags,
            default=tags,
            help="Tags help others discover your model",
        )

        # Model description
        st.markdown("#### Description")
        updated_description = st.text_area(
            "Provide a brief description of your model",
            value=description,
            help="This will appear on your model card and help others understand your model's purpose",
        )

        # Full model card content (for advanced users)
        st.markdown("#### Full Model Card (Markdown)")
        st.markdown(
            "Edit the full model card content if needed. This is in Markdown format."
        )
        updated_model_card = st.text_area(
            "Model Card Content", value=model_card_content, height=300
        )

        # Submit button
        submitted = st.form_submit_button(
            "Update Model Information", use_container_width=True
        )

        if submitted:
            # Update the model card
            with st.spinner("Updating model information..."):
                try:
                    repo_id = model_info.modelId

                    # If the user has edited the full model card, use that
                    if updated_model_card != model_card_content:
                        new_content = updated_model_card
                    else:
                        # Otherwise, update only tags and description in the existing card
                        # Update tags in YAML frontmatter
                        if yaml_match:
                            yaml_content = yaml_match.group(1)
                            if tags_match:
                                # Replace tags section
                                new_yaml = yaml_content.replace(
                                    tags_match.group(0),
                                    f"tags:\n"
                                    + "\n".join([f"- {tag}" for tag in selected_tags])
                                    + "\n",
                                )
                            else:
                                # Add tags section
                                new_yaml = (
                                    yaml_content
                                    + f"\ntags:\n"
                                    + "\n".join([f"- {tag}" for tag in selected_tags])
                                    + "\n"
                                )

                            new_content = model_card_content.replace(
                                yaml_match.group(0), f"---\n{new_yaml}---"
                            )
                        else:
                            # Add YAML frontmatter with tags
                            tags_yaml = (
                                "---\ntags:\n"
                                + "\n".join([f"- {tag}" for tag in selected_tags])
                                + "\n---\n\n"
                            )
                            new_content = tags_yaml + model_card_content

                        # Update description
                        if title_match and updated_description != description:
                            new_content = new_content.replace(
                                title_match.group(0),
                                title_match.group(0).replace(
                                    description, updated_description
                                ),
                            )

                    # Update the model card
                    success, _ = st.session_state.client.update_model_card(
                        repo_id, new_content
                    )

                    if success:
                        st.success("Model information updated successfully!")
                        # Refresh the models list
                        st.session_state.models = (
                            st.session_state.client.get_user_models()
                        )
                        return True, repo_id
                    else:
                        st.error("Failed to update model information")
                        return False, None

                except Exception as e:
                    st.error(f"Error updating model information: {str(e)}")
                    return False, None

    return False, None