package org.wikiwebserver.handler.http;

import java.awt.image.BufferedImage;
import java.io.*;
import java.util.Collection;
import java.util.Stack;
import java.util.zip.GZIPOutputStream;

import org.wikiwebserver.core.WareHouse;
import org.wikiwebserver.core.WikiCache;
import org.wikiwebserver.core.WikiCacheItem;
import org.wikiwebserver.core.WikiWebServer;
import org.wikiwebserver.handler.http.interfaces.CacheableHTTPResponse;
import org.wikiwebserver.handler.http.interfaces.HTTPResponder;
import org.wikiwebserver.handler.http.responder.BufferedImageResponder;
import org.wikiwebserver.handler.http.responder.ExceptionResponder;
import org.wikiwebserver.handler.http.responder.FileResponder;


public class HTTPWriter extends OutputStream {
    
    public HTTPWriter(HTTPHandler conn, OutputStream out) {
        this.conn = conn;
        this.rawOutputStream = out;
    }
    
    HTTPResponse startResponse() throws Exception {
    	
        this.response = new HTTPResponse();
        
        this.headersSent = false;
        this.contentExpected = true;
        
        // Assume OK until we know otherwise
        this.response.setCode(200);
        this.response.setInfo("OK");
        
        HTTPHeaders headers = new HTTPHeaders();
        headers.setDate("Date", System.currentTimeMillis());      
        
        headers.setAll(staticHeaders);       

        if (conn.getHTTPConfig() != null) {
            String serverId = conn.getHTTPConfig().staticIdentity("ServerID");
            headers.set("X-WikiWebServerID", serverId);
        }
                
        this.response.setHeaders(headers);    
        
        outputStreamStack.clear();
        outputStreamStack.push(rawOutputStream);
        
        return this.response;
    }
    
    private void writeString(String s) throws IOException { 
        write(s.getBytes("UTF8"));
    }     
    
    private void writeFile(File file) throws IOException {
        write(new FileResponder(file));        
    }
    
    private void writeImage(BufferedImage image) throws IOException {
        write(new BufferedImageResponder(image));        
    }      
    
    private void writeException(Throwable exception) throws IOException {
        write(new ExceptionResponder(exception));        
    }    
        
    private void runResponder(HTTPResponder responder) throws IOException {

        long time = System.currentTimeMillis();   
        HTTPHeaders responseHeaders = conn.getResponse().getHeaders();   
        
        WikiCacheItem cachedItem = null;        
        Object responseObject = null;
        boolean cacheable = false;
        long lastModifiedTime = time;
        long expireTime = time;
        
        if (responder instanceof CacheableHTTPResponse) {

            CacheableHTTPResponse cResponder = (CacheableHTTPResponse)responder;
            cResponder.init(conn);
            
            String method = conn.getRequest().getMethod();
            String key = cResponder.getCacheKey();
            expireTime = cResponder.getExpireTime();         
            
            cacheable = (key != null) && expireTime > time &&
                        "GET".equalsIgnoreCase(method) &&
                    	!"no-cache".equals(responseHeaders.getFirst("Cache-Control"));
            	
            	
            if (cacheable) {
                
                // Keep cacheable content specific to host, responder class and encoding        
                String host = conn.getRequest().getHeaders().getFirst("Host");                
                String resClass = responder.getClass().getName();

                key = host + "_" + resClass + "_" + key;
                
                cachedItem = cache.lookup(key);               
            } 

            if (cachedItem == null) {
                
                // A cache miss, generate response
                //System.out.println("Cache miss for URI: " + conn.getRequest().getUri());              
                
                response.setCacheKey(key);    
                response.setCacheExpireTime(expireTime);     
                
                responseObject = cResponder.respond(conn);          
                
                if (responseObject instanceof HTTPResponder) {
                    // The responder returned another responder             
                    runResponder((HTTPResponder)responseObject);
                    return;
                }
                
                // Invalidate cache if responder has made explicit request to avoid caching
                if ("no-cache".equals(responseHeaders.getFirst("Cache-Control"))) {
                    cacheable = false;
                }
            }
        }
        else {
            responseObject = responder.respond(conn);
            
            // A responder may return another responder
            if (responseObject instanceof HTTPResponder) {
                runResponder((HTTPResponder)responseObject);
                return;
            }             
        }      
        
        if (cacheable) {          
            
            if (cachedItem != null) {
                
                // Update this response with previous response headers
                responseHeaders.setDate("Expires", cachedItem.getExpireTime()); 
                
                long previousModified = cachedItem.getLastModifiedTime();
                if (previousModified > 0) {
                    responseHeaders.setDate("Last-Modified", previousModified);
                }
                
                long previousResponseTime = cachedItem.getCreatedTime();
                responseHeaders.set("Age", String.valueOf((time - previousResponseTime) / 1000));
                
                responseObject = cachedItem.getDataStream();
            }
            else {   
                responseHeaders.setDate("Expires", expireTime); 
            }
            
            responseHeaders.set("Pragma", "cache");
            responseHeaders.set("Cache-Control", "public");                  
        }       
        else {
            response.setCacheKey(null);
            response.setCacheExpireTime(0); 
            responseHeaders.setDate("Expires", time);
            responseHeaders.set("Pragma", "no-cache");          
            responseHeaders.set("Cache-Control", "no-cache"); 
        }
        
        
        if (responseHeaders.getFirst("Last-Modified") == null) {
            responseHeaders.setDate("Last-Modified", lastModifiedTime);
        }          

        if (cachedItem != null) {
            responseHeaders.set("ETag", "" + cachedItem.hashCode());
            // Pass on content type from cached response
            response.getHeaders().set("Content-Type", cachedItem.getContentType());            
        } else {
            responseHeaders.set("ETag", "" + response.hashCode());    
        }        
        
        if (responseObject != null) {
            
            if (responseObject instanceof String) {
                responseObject = ((String)responseObject).getBytes("utf8");
            }
            
            // Response object is the only response, specify length
            if (response.getNumBytesWritten() == 0) {
                if (responseObject instanceof byte[]) {
                    int length = ((byte[])responseObject).length; 
                    response.getHeaders().set("Content-Length", String.valueOf(length));
                }
                else if (responseObject instanceof ByteArrayOutputStream) {
                    int length = ((ByteArrayOutputStream)responseObject).size(); 
                    response.getHeaders().set("Content-Length", String.valueOf(length));
                }                
            }
            
            this.write(responseObject);
        }
        
        // No content
        if (response.getNumBytesWritten() == 0) {
        	response.getHeaders().set("Content-Length", String.valueOf(0));
            this.prepairStream();
        }
    }

    
    public void write(byte[] data) throws IOException { 
        write(data, 0, data.length);
    }    
    
    public void write(byte[] data, int start, int len) throws IOException { 
        
        prepairStream();
        
        if (!contentExpected) return;       
        
        if (response.isCacheable()) {
            response.getCacheStream().write(data, start, len);
        }
        
        OutputStream out = outputStreamStack.peek();
        out.write(data, start, len);
        
        this.response.incrementNumBytesWritten(len); 
        if (conn.getHTTPConfig() != null && 
            conn.getHTTPConfig().getHTTPMonitor() != null) {
            
            conn.getHTTPConfig().getHTTPMonitor().incrementNumBytesWritten(len);
        }
    }      
    
    
    public void write(Object obj) throws IOException {
        
        if (obj instanceof HTTPResponder) {
            // Dynamic response can (and will) vary based on user agent
            conn.getResponse().getHeaders().add("Vary", "User-Agent");           
            
            runResponder((HTTPResponder) obj);
        }         
        else if (obj instanceof ByteArrayOutputStream) {
            ((ByteArrayOutputStream) obj).writeTo(this);
        }          
        else if (obj instanceof byte[]) {
            write((byte[]) obj);
        }               
        else if (obj instanceof String) {
            writeString((String) obj);
        }        
        else if (obj instanceof File) {
            writeFile((File) obj);
        }          
        else if (obj instanceof BufferedImage) {
            writeImage((BufferedImage) obj);
        }
        else if (obj instanceof Throwable) {
            writeException((Throwable) obj);
        }        
        else if (obj == null) {
            // Ignore
        }
        else write(obj.toString());
    }      
        
    public void close() throws IOException {
        this.outputStreamStack.peek().close();
    }    
    
   
    public boolean cache() {
        
        // Cache the response for reuse
        if ((response.getCode() / 100 == 2) && response.isCacheable()) {     
            
            WikiCacheItem cacheItem = new WikiCacheItem();
            cacheItem.setLastModifiedTime(response.getHeaders().getFirstDate("Last-Modified"));
            cacheItem.setCreatedTime(response.getHeaders().getFirstDate("Date"));
            cacheItem.setContentType(response.getHeaders().getFirst("Content-Type"));
            cacheItem.setExpireTime(response.getHeaders().getFirstDate("Expires"));
            cacheItem.setDataStream(response.getCacheStream());
            
            cache.store(response.getCacheKey(), cacheItem);
            
            return true;
        }   
        return false;
    }

    public void finish() throws IOException {
        
        // Finish 'finishable' output streams bottom to top
        for (int i=outputStreamStack.size()-1; i>=0; i--) {
            OutputStream out = outputStreamStack.get(i);
            if (out instanceof ChunkedOutputStream) {
                ((ChunkedOutputStream) out).finish();
            }          
            else if (out instanceof GZIPOutputStream) {
                ((GZIPOutputStream) out).finish();
            }
        }
    }
    
    public void flush() throws IOException {
        this.outputStreamStack.peek().flush();
    }     
    
    
    /* START for output stream compatibility */
    
    public void write(int c) throws IOException {
        prepairStream();
        if (!contentExpected) return;
        
        if (response.isCacheable()) {
            response.getCacheStream().write(c);
        }
        
        OutputStream out = this.outputStreamStack.peek();
        out.write(c);
        
        this.response.incrementNumBytesWritten(1);
        
        if (conn.getHTTPConfig() != null && 
            conn.getHTTPConfig().getHTTPMonitor() != null) {
            
            conn.getHTTPConfig().getHTTPMonitor().incrementNumBytesWritten(1);
        }    
    }     
    
    /* END for output stream compatibility */
    
    
        
    public void prepairStream() throws IOException {
        
        if (this.headersSent) return;
        
        boolean applyChunking = false;
        boolean applyGZipping = false;
        
        HTTPHeaders requestHeaders = conn.getRequest().getHeaders();
        HTTPHeaders responseHeaders = response.getHeaders();
        
        // If no content type specified, assume HTML
        if (responseHeaders.get("Content-Type") == null) {
            responseHeaders.set("Content-Type", "text/html; charset=utf-8");  
        }
        
        // Process Connection header
        String connType = "close";
        if (conn.getReader().isActive() && conn.isReusable()) {
            int timeoutSecs = HTTPHandler.READ_TIMEOUT / 1000;
            connType = "Keep-Alive";
            responseHeaders.set("Keep-Alive", "timeout=" + timeoutSecs);
        }
        responseHeaders.set("Connection", connType);   
        
        String conLen = responseHeaders.getFirst("Content-Length");        
        
        boolean modified = true; // Assume modified    
        
        // Advanced cache checking
        Collection<String> ifNoneMatch = requestHeaders.get("If-None-Match");
        long ifModSince = requestHeaders.getFirstDate("If-Modified-Since");
        String responseETag = responseHeaders.getFirst("ETag");
        long lastModified = responseHeaders.getFirstDate("Last-Modified");
        if (ifNoneMatch != null && responseETag != null) { // Check for cache match
            responseETag = responseETag.trim();
            for (String testEtag : ifNoneMatch) {   
                if (testEtag.equals(responseETag)) {
                    modified = false;
                }
            }
        // Resort to basic cache checking
        } else if (ifModSince > 0 && lastModified > 0) {
            try {
                if (lastModified <= ifModSince) {
                    modified = false; // Old date - probably not modified
                }
            } catch (Exception ex) {
                // Thrown by invalid dates
                ex.printStackTrace(); 
            }
        }
        
        
        if (!modified) {
            this.response.setCode(304);
            this.response.setInfo("Not Modified");
            this.contentExpected = false;
        }          
        
        else if (conLen != null && Integer.parseInt(conLen) == 0) {
            this.response.setCode(204);
            this.response.setInfo("No content");
            this.contentExpected = false;
        } 
        
        else {

            // NB GZipping may already have been applied at this point
            if (isGZippingSuggested()) {                
                conLen = null;
                responseHeaders.remove("Content-Length"); // Not known until complete
                responseHeaders.set("Content-Encoding", "gzip");                
                applyGZipping = true;
            }    
            
            // Update the headers, GZipping breaks things
            if ("gzip".equalsIgnoreCase(responseHeaders.getFirst("Content-Encoding"))) {
                responseHeaders.set("Accept-Ranges", "none"); // Not supported                
                responseHeaders.add("Vary", "Accept-Encoding");  
            }
            
            // CHUNKING support
            applyChunking = connType.equalsIgnoreCase("Keep-Alive") &&
                          (conLen == null || Integer.parseInt(conLen) == -1);
            
            if (applyChunking) {
                int major = response.getVersionMajor();
                int minor = response.getVersionMinor();                
                if (major > 1 || (major == 1 && minor >= 1)) {
                    responseHeaders.set("Transfer-Encoding", "chunked");
                } else {
                    // Fall back to closing the connection upon completion
                    responseHeaders.set("Connection", "close");  
                    applyChunking = false;
                }
            }
        }
        
        //System.out.println(this.response.getCode() + " " + this.conn.getRequest().getUri());        
        
        // If only the head is requested, no content should be returned
        String method = conn.getRequest().getMethod().toUpperCase();    
        if ("HEAD".equals(method)) {
            this.contentExpected = false;
        }
        
        if (!this.contentExpected || this.response.getCode() == 204) {
            responseHeaders.remove("Transfer-Encoding");
            responseHeaders.remove("Content-Length");
        } 
        
        // System.out.println(response.toString());
        
        byte[] responseHeadData = response.toString().getBytes("US-ASCII");
        
        this.rawOutputStream.write(responseHeadData);
        this.response.incrementNumBytesWritten(responseHeadData.length);        
        this.headersSent = true;   
        
        if (conn.getHTTPConfig() != null && 
            conn.getHTTPConfig().getHTTPMonitor() != null) {
            
            conn.getHTTPConfig().getHTTPMonitor().incrementNumBytesWritten(responseHeadData.length);
        }        
        
        if (this.contentExpected) {
            if (applyChunking) {
                outputStreamStack.push(new ChunkedOutputStream(outputStreamStack.peek()));
            }        
            if (applyGZipping) {
                outputStreamStack.push(new GZIPOutputStream(outputStreamStack.peek(), GZIP_BUFFER_SIZE)); 
            }
        }
    }
    
    public boolean contentExpected() {
        return this.contentExpected;
    }
    
    private boolean isGZippingSuggested() {
        
        String accEnc = conn.getRequest().getHeaders().getFirst("Accept-Encoding");
        String range = conn.getRequest().getHeaders().getFirst("Range");
        String enc = response.getHeaders().getFirst("Content-Encoding");
        String type = response.getHeaders().getFirst("Content-Type");      
        
        // Supported?
        if (accEnc != null && accEnc.contains("gzip") && range == null) {
            // Suggested?
            if (enc != null || type == null) return false;
            
            return (type.startsWith("text/") ||
                    type.equals("audio/x-wav") ||
                    type.equals("audio/aiff"));
        }  
        
        return false;
    } 
    
    // TODO work out optimal buffer size
    private static final int GZIP_BUFFER_SIZE = 1024;
    
    private HTTPHandler conn;
    private OutputStream rawOutputStream;
    private Stack<OutputStream> outputStreamStack = new Stack<OutputStream>();
    
    private boolean headersSent;
    private boolean contentExpected;
    private HTTPResponse response; 
     
    
    private static WikiCache cache = WareHouse.getWikiCache();
    
    private static HTTPHeaders staticHeaders = new HTTPHeaders();
    static {
        String wwsVersion = WikiWebServer.VERSION_CODE;
        String javaVersion = System.getProperty("java.version");
        String os = System.getProperty("os.name") + " " +
                    System.getProperty("os.version") + " " +
                    System.getProperty("os.arch");
        int classLoaderId = Thread.currentThread().getContextClassLoader().hashCode();
        
        staticHeaders.set("Server", "WikiWebServer " + wwsVersion, "Java " + javaVersion, os);
        staticHeaders.set("Content-Language", "en");
        staticHeaders.set("X-Developer", "Michael Gardiner");
        staticHeaders.set("X-WikiWebServer-ClassLoaderID", String.valueOf(classLoaderId));        
        staticHeaders.set("X-Java-Source", "http://www.wikiwebserver.org");
        staticHeaders.set("Accept-Ranges", "none");        
        staticHeaders.setDate("X-Java-VMStartTime", WareHouse.getVMStartTime());        
    }
}