import gradio as gr
import json
import requests
import os
import pandas as pd
import folium
from geopy.geocoders import Nominatim
from geopy.exc import GeocoderTimedOut, GeocoderServiceError
import time
import random
from typing import List, Tuple, Optional
import io
import concurrent.futures
from tqdm import tqdm

# NuExtract API configuration
API_URL = "https://api-inference.huggingface.co/models/numind/NuExtract-1.5"
headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN', '')}"}

# Geocoding Service with improved performance
class GeocodingService:
    def __init__(self, user_agent: str = None, timeout: int = 10, rate_limit: float = 1.1):
        if user_agent is None:
            user_agent = f"python_geocoding_script_{random.randint(1000, 9999)}"

        self.geolocator = Nominatim(
            user_agent=user_agent,
            timeout=timeout
        )
        self.rate_limit = rate_limit
        self.last_request = 0
        self.cache = {}  # Simple in-memory cache for geocoding results

    def _rate_limit_wait(self):
        current_time = time.time()
        time_since_last = current_time - self.last_request
        if time_since_last < self.rate_limit:
            time.sleep(self.rate_limit - time_since_last)
        self.last_request = time.time()

    def geocode_location(self, location: str, max_retries: int = 3) -> Optional[Tuple[float, float]]:
        # Check cache first
        if location in self.cache:
            return self.cache[location]
            
        for attempt in range(max_retries):
            try:
                self._rate_limit_wait()
                location_data = self.geolocator.geocode(location)
                if location_data:
                    # Store in cache and return
                    self.cache[location] = (location_data.latitude, location_data.longitude)
                    return self.cache[location]
                # Cache None results too
                self.cache[location] = None
                return None
            except (GeocoderTimedOut, GeocoderServiceError) as e:
                if attempt == max_retries - 1:
                    print(f"Failed to geocode '{location}' after {max_retries} attempts: {e}")
                    self.cache[location] = None
                    return None
                time.sleep(2 ** attempt)  # Exponential backoff
            except Exception as e:
                print(f"Error geocoding '{location}': {e}")
                self.cache[location] = None
                return None
        return None

    def process_locations(self, locations: str, progress_callback=None) -> List[Optional[Tuple[float, float]]]:
        if pd.isna(locations) or not locations:
            return []

        # Handle special case with "dateline_locations" prefix
        if "dateline_locations" in locations:
            # Remove the prefix if present
            locations = locations.replace("dateline_locations", "").strip()
            
        # Improved location parsing to handle complex location names with commas
        # This regex-based approach attempts to identify well-formed location patterns
        try:
            import re
            
            # Try to find patterns like "City, Country" or standalone names
            # This handles cities like "Paris, France" as single entities
            location_pattern = re.compile(r'([A-Za-z\s]+(?:,\s*[A-Za-z\s]+)?)')
            matches = location_pattern.findall(locations)
            
            # Filter out empty matches and strip whitespace
            location_list = [match.strip() for match in matches if match.strip()]
            
            # If regex didn't work properly, fall back to a simpler approach
            if not location_list:
                # Simple space-based splitting as a fallback
                location_list = [loc.strip() for loc in locations.split() if loc.strip()]
                print(f"Using fallback location parsing: {location_list}")
                
        except Exception as e:
            print(f"Error parsing locations: {e}, using simple splitting")
            # Simple fallback if regex fails
            location_list = [loc.strip() for loc in locations.split() if loc.strip()]
            
        print(f"Parsed locations: {location_list}")
            
        # Process locations in parallel with a limited number of workers
        return self.process_locations_parallel(location_list, progress_callback)
        
    def process_locations_parallel(self, location_list, progress_callback=None, max_workers=4) -> List[Optional[Tuple[float, float]]]:
        """Process locations in parallel with progress tracking"""
        results = [None] * len(location_list)
        
        # Use a ThreadPoolExecutor for parallel processing
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit all tasks
            future_to_index = {executor.submit(self.geocode_location, loc): i 
                              for i, loc in enumerate(location_list)}
            
            # Process as they complete with progress updates
            total = len(future_to_index)
            completed = 0
            
            for future in concurrent.futures.as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    results[index] = future.result()
                except Exception as e:
                    print(f"Error processing location: {e}")
                    results[index] = None
                
                # Update progress
                completed += 1
                if progress_callback:
                    progress_callback(completed, total)
                else:
                    print(f"Geocoded {completed}/{total} locations")
                    
        return results

# Mapping Functions
def create_location_map(df: pd.DataFrame,
                       coordinates_col: str = 'coordinates',
                       places_col: str = 'places',
                       title_col: Optional[str] = None) -> folium.Map:
    # Initialize the map
    m = folium.Map(location=[0, 0], zoom_start=2)
    all_coords = []

    # Process each row in the DataFrame
    for idx, row in df.iterrows():
        coordinates = row[coordinates_col]
        places = row[places_col].split(',') if pd.notna(row[places_col]) else []
        title = row[title_col] if title_col and pd.notna(row[title_col]) else None

        # Skip if no coordinates
        if not coordinates:
            continue

        # Make sure places and coordinates lists have the same length
        # If places list is shorter, pad it with unnamed locations
        while len(places) < len(coordinates):
            places.append(f"Unnamed Location {len(places)+1}")
            
        # Add individual markers for each location
        for i, coord in enumerate(coordinates):
            if coord is not None:  # Skip None coordinates
                lat, lon = coord
                # Safely get place name, use a default if index is out of range
                place_name = places[i].strip() if i < len(places) else f"Location {i+1}"

                # Create popup content
                popup_content = f"<b>{place_name}</b>"
                if title:
                    popup_content += f"<br>{title}"

                # Add marker to the map
                folium.Marker(
                    location=[lat, lon],
                    popup=folium.Popup(popup_content, max_width=300),
                    tooltip=place_name,
                ).add_to(m)

                all_coords.append([lat, lon])

    # If we have coordinates, fit the map bounds to include all points
    if all_coords:
        m.fit_bounds(all_coords)

    return m

# Processing Functions with progress updates
def process_excel(file, places_column, progress=None):
    # Check if file is None
    if file is None:
        return None, "No file uploaded", None
    
    try:
        # Update progress
        if progress:
            progress(0.1, "Reading Excel file...")
        
        # Handle various file object types that Gradio might provide
        if hasattr(file, 'name'):
            # Gradio file object
            df = pd.read_excel(file.name)
        elif isinstance(file, bytes):
            # Raw bytes
            df = pd.read_excel(io.BytesIO(file))
        else:
            # Assume it's a filepath string
            df = pd.read_excel(file)
        
        if places_column not in df.columns:
            return None, f"Column '{places_column}' not found in the Excel file. Available columns: {', '.join(df.columns)}", None
        
        # Print column names and first few rows for debugging
        print(f"Columns in Excel file: {df.columns.tolist()}")
        print(f"First 3 rows of data:\n{df.head(3)}")
        
        if progress:
            progress(0.2, "Initializing geocoding...")
        
        # Initialize the geocoding service
        geocoder = GeocodingService(user_agent="gradio_map_visualization_app")
        
        # Function to update progress during geocoding
        def geocoding_progress(completed, total):
            if progress:
                # Scale progress between 20% and 80%
                progress_value = 0.2 + (0.6 * (completed / total))
                progress(progress_value, f"Geocoding {completed}/{total} locations...")
        
        # Process locations and add coordinates with progress tracking
        print("Starting geocoding process...")
        
        # Process each row with progress updates
        coordinates_list = []
        total_rows = len(df)
        
        # Create a helper function to safely parse location data from each row
        def parse_excel_locations(location_data):
            """Safely parse location data from Excel cell"""
            if pd.isna(location_data):
                return []
                
            # Convert to string to handle numeric or other data types
            location_data = str(location_data).strip()
            
            # Skip empty strings
            if not location_data:
                return []
                
            # Look for recognized patterns and split accordingly
            # First, check if it's a comma-separated list
            if "," in location_data:
                # This could be a list like "Berlin, Hamburg, Munich"
                # Or it could contain locations like "Paris, France"
                
                # Try to intelligently parse based on common patterns
                try:
                    import re
                    
                    # Pattern to match city-country pairs or standalone names
                    # Examples: "Paris, France" or "Berlin" or "New York, NY, USA"
                    location_pattern = re.compile(r'([A-Za-z\s]+(?:,\s*[A-Za-z\s]+){0,2})')
                    matches = location_pattern.findall(location_data)
                    
                    locations = [match.strip() for match in matches if match.strip()]
                    
                    # If our pattern matching didn't work, fall back to simple comma splitting
                    if not locations:
                        locations = [loc.strip() for loc in location_data.split(',') if loc.strip()]
                        
                    return locations
                    
                except Exception as e:
                    print(f"Regex parsing failed: {e}")
                    # Fallback to simple comma splitting
                    return [loc.strip() for loc in location_data.split(',') if loc.strip()]
            
            # Otherwise, treat it as a single location or space-separated list
            else:
                # Check if it might be space-separated
                potential_locations = location_data.split()
                
                # If it just looks like one word with no spaces, return it as a single location
                if len(potential_locations) == 1:
                    return [location_data]
                    
                # If it has multiple words, it could be a single location name with spaces
                # or multiple space-separated locations
                # For safety, treat it as a single location
                return [location_data]
        
        for idx, row in df.iterrows():
            location_data = row[places_column]
            print(f"Processing row {idx+1}/{total_rows}, location data: {location_data}")
            
            # Parse the locations from the Excel cell
            location_list = parse_excel_locations(location_data)
            print(f"Parsed locations: {location_list}")
            
            # Now geocode each location
            coords = []
            for location in location_list:
                coord = geocoder.geocode_location(location)
                coords.append(coord)
                # Update progress
                if progress_callback:
                    progress_callback(len(coords), len(location_list))
            
            coordinates_list.append(coords)
            print(f"Processed row {idx+1}/{total_rows}, found coordinates: {coords}")
            
        df['coordinates'] = coordinates_list
        
        if progress:
            progress(0.8, "Creating map...")
        
        # Create the map
        map_obj = create_location_map(df, coordinates_col='coordinates', places_col=places_column)
        
        # Save the map to a temporary HTML file
        temp_map_path = "temp_map.html"
        map_obj.save(temp_map_path)
        
        # Save the processed DataFrame to Excel
        if progress:
            progress(0.9, "Saving results...")
            
        processed_file_path = "processed_data.xlsx"
        df.to_excel(processed_file_path, index=False)
        
        # Statistics
        total_locations = len(df)
        successful_geocodes = sum(1 for coords in coordinates_list for coord in coords if coord is not None)
        failed_geocodes = sum(1 for coords in coordinates_list for coord in coords if coord is None)
        
        stats = f"Total data rows: {total_locations}\n"
        stats += f"Successfully geocoded locations: {successful_geocodes}\n"
        stats += f"Failed to geocode locations: {failed_geocodes}"
        
        if progress:
            progress(1.0, "Processing complete!")
            
        return temp_map_path, stats, processed_file_path
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"Error processing Excel file: {error_details}")
        
        if progress:
            progress(1.0, f"Error: {str(e)}")
        return None, f"Error processing file: {str(e)}\n\nDetails: {error_details}", None

# NuExtract Functions
def extract_info(template, text):
    try:
        # Format prompt according to NuExtract-1.5 requirements
        prompt = f"<|input|>\n### Template:\n{template}\n### Text:\n{text}\n\n<|output|>"
        
        # Call API
        payload = {
            "inputs": prompt,
            "parameters": {
                "max_new_tokens": 1000,
                "do_sample": False
            }
        }
        
        response = requests.post(API_URL, headers=headers, json=payload)
        
        # If the model is loading, inform the user
        if response.status_code == 503:
            response_json = response.json()
            if "error" in response_json and "loading" in response_json["error"]:
                estimated_time = response_json.get("estimated_time", "unknown")
                return f"⏳ Model is loading (ETA: {int(float(estimated_time)) if isinstance(estimated_time, (int, float, str)) else 'unknown'} seconds)", "Please try again in a few minutes"
        
        if response.status_code != 200:
            return f"❌ API Error: {response.status_code}", response.text
        
        # Process result
        result = response.json()
        
        # Handle different response formats with careful error handling
        try:
            if isinstance(result, list):
                if len(result) > 0:
                    result_text = result[0].get("generated_text", "")
                else:
                    return "❌ Empty result list from API", "{}"
            else:
                result_text = str(result)
            
            # Split at output marker if present
            if "<|output|>" in result_text:
                split_parts = result_text.split("<|output|>")
                if len(split_parts) > 1:
                    json_text = split_parts[1].strip()
                else:
                    json_text = result_text  # Fallback if split didn't work as expected
            else:
                json_text = result_text
            
            # Try to parse as JSON
            try:
                extracted = json.loads(json_text)
                formatted = json.dumps(extracted, indent=2)
            except json.JSONDecodeError:
                return "❌ JSON parsing error", json_text
                
            return "✅ Success", formatted
        except Exception as inner_e:
            return f"❌ Error processing API result: {str(inner_e)}", "{}"
    except Exception as e:
        return f"❌ Error: {str(e)}", "{}"

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Historical Data Analysis Tools")
    
    with gr.Tabs():
        with gr.TabItem("Text Extraction"):
            gr.Markdown("## NuExtract-1.5 Structured Data Extraction")
            
            with gr.Row():
                with gr.Column():
                    template = gr.Textbox(
                        label="JSON Template", 
                        value='{"earthquake location": "", "dateline location": ""}',
                        lines=5
                    )
                    text = gr.Textbox(
                        label="Text to Extract From",
                        value="Neues Erdbeben in Japan. Aus Tokio wird berichtet, daß in Yokohama bei einem Erdbeben sechs Personen getötet und 22 verwundet, in Tokio vier getötet und 22 verwundet wurden. In Yokohama seien 6VV Häuser zerstört worden. Die telephonische und telegraphische Verbindung zwischen Tokio und Osaka ist unterbrochen worden. Der Trambahnverkehr in Tokio liegt still. Auch der Eisenbahnverkehr zwischen Tokio und Yokohama ist unterbrochen. In Sngamo, einer Vorstadt von Tokio sind Brände ausgebrochen. Ein Eisenbahnzug stürzte in den Vajugawafluß zwischen Gotemba und Tokio. Sechs Züge wurden umgeworfen. Mit dem letzten japanischen Erdbeben sind seit eineinhalb Jahrtausenden bis heute in Japan 229 größere Erdbeben zu verzeichnen gewesen.",
                        lines=8
                    )
                    extract_btn = gr.Button("Extract Information", variant="primary")
                
                with gr.Column():
                    status = gr.Textbox(label="Status")
                    output = gr.Textbox(label="Output", lines=10)
            
            extract_btn.click(
                fn=extract_info,
                inputs=[template, text],
                outputs=[status, output]
            )
        
        with gr.TabItem("Geocoding & Mapping"):
            gr.Markdown("## Location Mapping Tool")
            
            with gr.Row():
                with gr.Column():
                    excel_file = gr.File(label="Upload Excel File")
                    places_column = gr.Textbox(label="Places Column Name", value="places")
                    process_btn = gr.Button("Process and Map", variant="primary")
                
                with gr.Column():
                    progress_bar = gr.Progress()
                    map_output = gr.HTML(label="Map Visualization")
                    stats_output = gr.Textbox(label="Statistics", lines=3)
                    processed_file = gr.File(label="Processed Data", visible=True, interactive=False)
            
            def process_and_map(file, column, progress=gr.Progress()):
                if file is None:
                    return None, "Please upload an Excel file", None
                
                try:
                    # Initialize progress
                    progress(0, "Starting process...")
                    
                    # Process the file with progress updates
                    map_path, stats, processed_path = process_excel(file, column, progress)
                    
                    if map_path and processed_path:
                        with open(map_path, "r") as f:
                            map_html = f.read()
                        
                        return map_html, stats, processed_path
                    else:
                        return None, stats, None
                except Exception as e:
                    return None, f"Error: {str(e)}", None
            
            process_btn.click(
                fn=process_and_map,
                inputs=[excel_file, places_column],
                outputs=[map_output, stats_output, processed_file]
            )

if __name__ == "__main__":
    demo.launch()